from PYEVALB import scorer
import os
import string
import trees


def load_lines(fl):
    with open(fl,'r') as f:
        lines = [line.strip() for line in f.readlines()]
    return lines

def tree2str(tree):
    return ' '.join(''.join(ch for ch in tree if not ch.isupper()).replace('(','').replace(')','').translate(str.maketrans('', '', string.punctuation)).strip().split())


exp = ['0','1','2']



model2output = {}
model2scores = {}
# load all the files into model2output, as lists of lines

out = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/code/self_attn_speech_parser/output/turn_opensmile'
data = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/turn_opensmile'

gold_file = os.path.join(data,'turn_dev.trees')
gold_lines = load_lines(gold_file)

for ex in exp:
    seg_model = 'sp_seg_dev=89.84.pt'
    #seg_model = 'nonsp_seg_dev=87.60.pt'
    parser_training = 'silver'
    file_name = os.path.join(out,f'{seg_model}.dev.beam3.pred.{ex}.{parser_training}.parsed.{parser_training}.turns')
    lines = load_lines(file_name)
    scores = [[float(score) for score in line.split()] for line in load_lines(file_name+'.scores')]
    model2output[ex] = lines
    model2scores[ex] = scores
    assert len(lines) == len(gold_lines)
    for i,(test,gold) in enumerate(zip(lines,gold_lines)):
        #print(ex)
        #print(i)
        #print(tree2str(test))
        #print(tree2str(gold))
        #print('-'*50)
        try:
            assert tree2str(test)==tree2str(gold)
        except:
            import pdb;pdb.set_trace()
    
gold_ids = load_lines(os.path.join(data,'turn_dev.ids'))
best_trees = []

for i,idnum in enumerate(gold_ids):
    best_score = -100000
    best_tree = ''
    for ex in exp:
        scores = model2scores[ex][i]
        score = sum(scores)#/len(scores)
        if score > best_score:
            best_tree = model2output[ex][i]
            best_score = score
    best_trees.append(best_tree)
        

best_file = file_name.replace('.{ex}.{parser_training}.parsed.{parser_training}.turns','.turns.best')
with open(best_file,'w') as f:
    for tree in best_trees:
        f.write(tree)
        f.write('\n')

assert len(gold_ids)==len(best_trees)

eval_file = best_file+'.eval'
scr = scorer.Scorer()
scr.evalb(gold_file,best_file,eval_file)
