from misc.find_all_hard_negatives import *
import torch
import pickle
from matplotlib import pyplot as plt
from datetime import datetime

def print_time():
    now = datetime.now()
    current_time = now.strftime("%H:%M:%S")
    print(current_time)
    return

dataset = "wn18rr"
mode = "COUNT"
datapath="data/"+dataset+"/train.txt"
filedatapath = dataset+"-all_hard_negatives_dict-"+mode+"-factorized.pkl"
entitypath = "data/"+dataset+"/entities.dict"

df = pd.read_csv(datapath, sep="\t", names=["head", "relation", "tail"], dtype=str)
print("dataframe loaded...", datapath)

#entities = list(set(df["head"].tolist() + df["tail"].tolist()))
entities = pd.read_csv(entitypath, sep="\t", header=None, dtype=str)[1].values

factorized_entities = pd.factorize(entities)
print("factorized_entities created")


ent2id = {}
id2ent = {}
for x in range(len(entities)):
  id2ent[x] = factorized_entities[1][x]
  ent2id[factorized_entities[1][x]] = x


df["head_ind"] = df["head"].apply(lambda x: ent2id[x])
df["tail_ind"] = df["tail"].apply(lambda x: ent2id[x])


factorized_relations = pd.factorize(df["relation"])
tensor_heads = torch.tensor(df["head_ind"])
tensor_tails = torch.tensor(df["tail_ind"])
tensor_relations = torch.tensor(factorized_relations[0])
tensor_triples = torch.stack([tensor_heads, tensor_relations, tensor_tails], dim=1)

print("triples are converted in a tensor form")
print_time()
neighbor_sets = find_neighbor_sets(tensor_triples, factorized_entities)
print("neighboring sets are calculated")
print_time()
all_hard_negatives_t = find_all_hard_negatives_per_entity(neighbor_sets,factorized_entities[0], mode=mode.lower())
print("all_hard_negatives are created")
print_time()


with open(filedatapath, 'wb') as f:
    pickle.dump(all_hard_negatives_t, f)
print("all_hard_negatives are dumped...")

df = pd.DataFrame(list(zip(factorized_entities[0], factorized_entities[1])))
df.to_csv("fact_entities.dict", index=False, header=False, sep="\t")
print("fact_entities is stored")

#--------------------
#For analysis of all candidate hard negatives

print("------------------")
print("Now analysis begins...")

file = open(filedatapath,'rb')
all_hard_negatives = pickle.load(file)
print("all_hard_negatives are loaded as dict")

total_sum = []
for x in all_hard_negatives.items():
  total_sum.append(len(x[1]))
print("Average candidate hard_negative per entity", str(sum(total_sum)/len(all_hard_negatives)))


total_sum_df = pd.DataFrame(total_sum)
total_sum_df.describe()

stats_df = total_sum_df \
.groupby(0) \
[0] \
.agg('count') \
.pipe(pd.DataFrame) \
.rename(columns = {0: 'frequency'})

# PDF
stats_df['pdf'] = stats_df['frequency'] / sum(stats_df['frequency'])

# CDF
stats_df['cdf'] = stats_df['pdf'].cumsum()
stats_df = stats_df.reset_index()
stats_df

stats_df_plot = stats_df.plot(x = 0, y = ['pdf', 'cdf'])
fig = stats_df_plot.get_figure()
fig.savefig(dataset+"-stats_df_plot.png")


plt.plot(total_sum)
