import pickle as pkl
import re

pkl_path = '../preprocess_input/test_gcn_gtr_lstm.pkl'  # test only
with open(pkl_path,'rb') as f:
    data = pkl.load(f)


pred_file = 'gcn_gtr_lstm_augment_original.txt'
with open(pred_file,'r') as f:
    pred_lines = f.readlines()

new_lines=[]
for i,line in enumerate(pred_lines):
    ex = data[i][0]
    dct={}
    for tri in ex:
        dct[tri[0][1]] = tri[0][0]
        dct[tri[1][1]] = tri[1][0]
        dct[tri[2][1]] = tri[2][0]

    for k,v in dct.items():
        line = line.replace(k.lower(), v.lower())
    line = ' '.join(re.split('(\W)',line)).replace('    ',' ').replace('   ',' ').replace('  ',' ').strip()
    new_lines.append(line.lower())

    # stop=1

pred_file = pred_file.split('.')[0]+'_relex.txt'
with open(pred_file,'w+') as f:
    f.write('\n'.join(new_lines))

stop=1