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','3','4']

model2output = {}
# 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'

for ex in exp:
    #file_name = os.path.join(out,f'sp_seg_dev=89.84.pt.dev.beam5.pred.{ex}.gold.parsed.gold.turns')
    file_name = os.path.join(out,f'sp_seg_dev=89.84.pt.dev.beam5.pred.{ex}.silver.parsed.silver.turns')
    lines = load_lines(file_name)
    model2output[ex] = lines


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

gold_trees,turn_ids =  trees.load_trees_with_idx(gold_file, gold_ids)

gold_treestr = [tr.linearize() for tr in gold_trees]

id2goldtree = dict(zip(turn_ids,gold_treestr))

id2predtree = {'0':{},
               '1':{},
               '2':{},
               '3':{},
               '4':{}}

scr = scorer.Scorer()

id2matched = {'0':{},
              '1':{},
              '2':{},
              '3':{},
              '4':{}}

id2gold = {'0':{},
           '1':{},
           '2':{},
           '3':{},
           '4':{}}

id2test = {'0':{},
           '1':{},
           '2':{},
           '3':{},
           '4':{}}


model2outputdict = {'0':{},
                    '1':{},
                    '2':{},
                    '3':{},
                    '4':{}}


dev_subset_file = '/afs/inf.ed.ac.uk/group/project/prosody/asr_preps/ranking_split_turns.json'
import json
dev_subset_ids = set(json.load(open(dev_subset_file,'r'))['dev'])

#turn_ids = [idnum for idnum in turn_ids if idnum in dev_subset_ids]



print('Store bracket scores ...')
for model in exp:
    test_lines = model2output[model]

    for i,turn in enumerate(turn_ids):#dev_subset_ids):
        #for i,turn in enumerate(turn_ids):

        gold_line = gold_lines[i]
        test_line = test_lines[i]
        id2predtree[model][turn] = test_line
        model2outputdict[model][turn] = test_line
        assert tree2str(gold_line)==tree2str(test_line)
        mat,gol,tes = scr.get_bracket_counts_from_tree(gold_line,test_line)
        id2matched[model][turn] = mat
        id2gold[model][turn] = gol
        id2test[model][turn] = tes

out_tmp = 'oracle_rerank_dict.json'
out_dict = {'id2matched':id2matched,
            'id2gold':id2gold,
            'id2test':id2test,
            'id2goldtree':id2goldtree}

with open(out_tmp,'w') as f:
    json.dump(out_dict,f)

def fscore(mat,gol,tes):
    rec = 100*mat/gol
    prec = 100*mat/tes
    if rec == 0 and prec == 0:
        return 0
    else:
        fs = (2*prec*rec)/(prec+rec)
        return fs

matches = []
golds = []
tests = []

best_trees = []
best_models = []
for idnum in dev_subset_ids:
    #for idnum in turn_ids:

    # calculate the f score for each model, pick the highest, save its mat/gol/tes values in a list
    best_fscore = -1000000
    best_ex = None
    for ex in exp:
        fs = fscore(id2matched[ex][idnum],id2gold[ex][idnum],id2test[ex][idnum])
        if fs > best_fscore:
            best_fscore = fs
            best_ex = ex
            mat = id2matched[ex][idnum]
            gol = id2gold[ex][idnum]
            tes = id2test[ex][idnum]
    best_models.append(best_ex)
    best_trees.append(model2outputdict[best_ex][idnum])
    matches.append(mat)
    golds.append(gol)
    tests.append(tes)

total_matches = sum(matches)
total_golds = sum(golds)
total_tests = sum(tests)

print(f'f score: {fscore(total_matches,total_golds,total_tests)}')

out_trees = file_name+'.best.trees'
out_models = file_name+'.best.n'
out_gold = file_name+'.subset'

with open(out_trees,'w') as f_tree:
    with open(out_models,'w') as f_mod:
        with open(out_gold,'w') as f_gold:
            for idnum,tree,ex in zip(dev_subset_ids,best_trees,best_models):
                f_tree.write(f'{tree}\n')
                f_mod.write(f'{ex}\n')
                f_gold.write(f'{id2goldtree[idnum]}\n')


scr.evalb(out_gold,out_trees,out)
