import json
from tools.framenet.naive_identifier import FrameIdentifier

test_file_path = '/home/gqin2/data/framenet/full/test.jsonl'
test_sentences = [
    json.loads(line) for line in open(test_file_path)
]
test_set = []
for ann in test_sentences:
    for fr in ann['frame']:
        test_set.append((fr['name'], ann['text'][fr['target'][0]: fr['target'][-1]+1], fr['lu']))

fi = FrameIdentifier()


tp = fp = fn = 0
fails = []
for frame, target_words, lu in test_set:
    pred = fi(target_words)
    if frame in pred:
        tp += 1
        fp += len(pred) - 1
    else:
        fp += len(pred)
        fn += 1
        fails.append((frame, target_words, pred, lu))

fails.sort(key=lambda x: x[0])
for frame, target_words, pred, lu in fails:
    print(frame, ' '.join(target_words), ' '.join(pred), lu, sep='\t')

print(f'tp={tp}, fp={fp}, fn={fn}')
print(f'precision={tp/(tp+fp)}')
print(f'recall={tp/(tp+fn)}')