import pickle
import numpy as np
import pandas as pd
import re

root = '../conceptnet/data/cpnet/'
concept_csv = pd.read_csv(root+'conceptnet.en.csv', 
                            sep='\t', 
                            names=['relation', 'head', 'tail', 'weight'],
                            header=None)
with open(root+'concept.txt', 'r') as f:
    entity_dic = {line.strip('\n'): i for i, line in enumerate(f)}
with open(root+'relation.txt', 'r') as f:
    relation_dic = {line.strip('\n'): i for i, line in enumerate(f)}
entity_embeds = np.load('../conceptnet/data/transe/glove.transe.sgd.ent.npy')
relation_embeds = np.load('../conceptnet/data/transe/glove.transe.sgd.ent.npy')

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--t', dest='token')
args = parser.parse_args()

df = concept_csv[concept_csv['head'].str.contains(args.token, na=False)]
# df = concept_csv[concept_csv['head']==args.token]
kdf = df.sort_values(by=['weight'])[-5:]
print(kdf)



