import numpy as np
import scipy.sparse as sp
import torch
import time
import random

from utils.tool import read_data, write_dic, dictionary, normalize, sparse_mx_to_torch_sparse_tensor
    

def encoding_test(test_graph_path, test_fact_path, train_dataset = "fb237_v1"):
    
    """load test-graph and test-facts, and do the encoding on the test-graph"""
     
    
    t_start = time.time()
    
    path = "data"



    #these two paths are for loading
    relation_dic_path = "{}/{}/train/relation-dic.txt".format(path, train_dataset)
    type_dic_path = "{}/{}/train/type-dic.txt".format(path, train_dataset)

    test_graph_triples = read_data(test_graph_path)
    test_fact_triples_with_label = read_data(test_fact_path)
    
    #load relation dic and type dic generated by training 
    f_relation_dic = open(relation_dic_path)
    relations = []
    for line in f_relation_dic:
        relation_new = line.strip().split("\t")[1]
        relations.append(relation_new)
    f_type_dic = open(type_dic_path)
    types = []
    for line in f_type_dic:
        type_new = line.strip().split("\t")[1]
        types.append(type_new)    

    relation_set = set(relations)

    all_triples_with_label = test_graph_triples + test_fact_triples_with_label

            
    test_graph_real_triples = []
    test_graph_type_triples = []
    for triple in test_graph_triples:
        if triple[1] != "<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>":
            test_graph_real_triples.append(triple)
        else:
            test_graph_type_triples.append(triple)

    test_fact_real_triples_with_label = []
    test_fact_type_triples_with_label = []
    for triple in test_fact_triples_with_label:
        if triple[1] != "<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>":
            test_fact_real_triples_with_label.append(triple)
        else:
            test_fact_type_triples_with_label.append(triple)
    
    all_real_triples_with_label = []
    all_type_triples_with_label = []
    
    constant_set = set()
    for triple in all_triples_with_label:
        if triple[1] != "<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>":
            constant_set.add(triple[0])
            constant_set.add(triple[2])
            all_real_triples_with_label.append(triple)
        else:
            constant_set.add(triple[0])
            all_type_triples_with_label.append(triple)
    
    constants = list(constant_set)

    



    constant2index = dictionary(constants)
    relation2index = dictionary(relations)
    type2index = dictionary(types)
    #print("time:",time.time()-t_start)
    #generate list of pairs for encoding
    
    pairs = []
    pair_set = set()
    for triple in all_real_triples_with_label:
        sub_idx = constant2index[triple[0]]
        obj_idx = constant2index[triple[2]]
        if sub_idx < obj_idx:
            if (sub_idx, obj_idx) not in pair_set:
                pair_set.add((sub_idx, obj_idx))
                pairs.append((sub_idx, obj_idx))
        if sub_idx > obj_idx:
            if (obj_idx, sub_idx) not in pair_set:
                pair_set.add((obj_idx, sub_idx))
                pairs.append((obj_idx, sub_idx))

    
    for constant_idx in range(len(constants)):
        pairs.append((constant_idx, constant_idx))
        pair_set.add((constant_idx, constant_idx))
    pair2index = dictionary(pairs)


    
    s_time = time.time()
    
    #collect related pairs for each constant
    
    pairs_for_constant = dict([(i,set()) for i in range(len(constants))])
    p_idx = 0
    for pair in pairs:
        p_idx = pair2index[pair]
        c1 = pair[0]
        c2 = pair[1]
        pairs_for_constant[c1].add(p_idx)
        pairs_for_constant[c2].add(p_idx)
    
    #collect neighbors for each pair node
    
    pneighbors_for_pair = dict([(i,set()) for i in range(len(pairs))])
    for c_idx in range(len(constants)):
        pairs_c = set(pairs_for_constant[c_idx])
        #pair and n_pair would contain one common constant
        for pair in pairs_c:
            for n_pair in pairs_c:
                if pair != n_pair:
                    pneighbors_for_pair[pair].add(n_pair)     
                    
    #generate edge list
    
    edges = []
    
    for i in range(len(pairs)):
        pneighbors = pneighbors_for_pair[i]
        for pneighbor in pneighbors:
            edges.append([i, pneighbor])
            edges.append([pneighbor, i])
            
    #print("Finished generating edges", time.time() - s_time)
    
    #generate a normalized adjencency matrix (strategy for GCN)
    #print(edges)
    edges = np.array(edges)
    adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])), shape=(len(pairs), len(pairs)), dtype=np.float32)
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)
    
    adj = normalize(adj + sp.eye(adj.shape[0]))
    adj = sparse_mx_to_torch_sparse_tensor(adj)
    del edges
    #print("Total time for adj: {:.4f}s".format(time.time() - s_time))
    
    #print("Start to generate features, labels, and masks")
    
    def initialize(test_graph_real_triples, test_graph_type_triples, test_fact_real_triples_with_label, test_fact_type_triples_with_label):
    
        labels = torch.zeros(len(pairs), len(types) + 2*len(relations))
        masks = torch.zeros(len(pairs),  len(types) + 2*len(relations))
        features = torch.zeros(len(pairs),  len(types) + 2*len(relations))
        
        #labels and masks are generated for all triples in test-facts (pos&neg)
        
        for triple in test_fact_type_triples_with_label:
            cons = triple[0]
            typ = triple[2]
            label = triple[3]
            pair_idx= pair2index[(constant2index[cons], constant2index[cons])]
            typ_idx = type2index[typ]
            if label == "1":
                labels[pair_idx][typ_idx] = 1
            elif label == "0":
                labels[pair_idx][typ_idx] = 0
            masks[pair_idx][typ_idx] = 1 

        for triple in test_fact_real_triples_with_label:
            sub = triple[0]
            rel = triple[1]
            obj = triple[2]
            label = triple[3]
            sub_idx = constant2index[sub]
            rel_idx = relation2index[rel]
            obj_idx = constant2index[obj]

            try:
                pair_idx = pair2index[(sub_idx, obj_idx)]        
            except:
                pair_idx = pair2index[(obj_idx, sub_idx)]
                rel_idx = rel_idx + len(relations)
            if label == "1":
                labels[pair_idx][len(types) + rel_idx] = 1
            elif label == "0":
                labels[pair_idx][len(types) + rel_idx] = 0
            masks[pair_idx][len(types) + rel_idx] = 1 
        
        #features are generated for all triples in test-graph (pos&neg)

        for triple in test_graph_type_triples:
            cons = triple[0]
            typ = triple[2]
            pair_idx= pair2index[(constant2index[cons], constant2index[cons])]
            typ_idx = type2index[typ]
            features[pair_idx][typ_idx] = 1



        for triple in test_graph_real_triples:
            sub = triple[0]
            rel = triple[1]
            obj = triple[2]
            sub_idx = constant2index[sub]
            rel_idx = relation2index[rel]
            obj_idx = constant2index[obj]

            try:
                pair_idx = pair2index[(sub_idx, obj_idx)]        
            except:
                pair_idx = pair2index[(obj_idx, sub_idx)]
                rel_idx = rel_idx + len(relations)
            features[pair_idx][len(types) + rel_idx] = 1
            


        features.requires_grad = True
        labels.requires_grad = False

        return features, labels, masks
    
    features, labels, masks = initialize(test_graph_real_triples, test_graph_type_triples, test_fact_real_triples_with_label, test_fact_type_triples_with_label)
    
    num_type = len(types)
    num_relation = len(relations)
    def triple2index(triple_now):
        sub_idx = constant2index[triple_now[0]]
        try:
            relation_idx = relation2index[triple_now[1]]
        except:
            pair_idx = pair2index[(sub_idx, sub_idx)]
            dim_idx = type2index[triple_now[2]]
            return pair_idx, dim_idx
        
        obj_idx = constant2index[triple_now[2]]
        if (sub_idx, obj_idx) in pair_set:
            pair_idx = pair2index[(sub_idx, obj_idx)]
            dim_idx = len(types) + relation_idx
        elif (obj_idx, sub_idx) in pair_set:
            pair_idx = pair2index[(obj_idx, sub_idx)]
            dim_idx = len(types) + len(relations) + relation_idx
        else:
            print(triple_now, sub_idx, relation_idx, obj_idx)
            print("wrong")
        return pair_idx, dim_idx
    hits_true = []
    for triple in test_fact_triples_with_label:
        if triple[-1] == "1":
            hits_true.append(triple2index(triple))
    #print("Finished generation")
    
    #print("Total time elapsed for encoding: {:.4f}s".format(time.time() - t_start))

            
    return adj, features, labels, masks, num_type, num_relation, constants, relations, types, pairs, hits_true

