import trees
import os

data_config = 'turn_part_fixed'

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

#predf = os.path.join(out,'turn_sp_parse_72240_dummy_penalty.5_dev=86.36.pt_dev_predicted.txt')
predf = os.path.join(out,'turn_sp_glove_72240_part_fixed_dev=86.34.pt_dev_predicted.txt')
goldf = os.path.join(data, 'turn_dev.trees')
idf = os.path.join(data,'turn_dev.ids')

pred_total = 0
gold_total = 0

preds,_ = trees.load_trees_with_idx(predf,idf,strip_top=True)
golds,_ = trees.load_trees_with_idx(goldf,idf,strip_top=True)

def count_nodes(tree):
    counter = 0
    node_queue = [tree]
    while node_queue:

        for child in node_queue[0].children:
            if not child.is_leaf():
                node_queue.append(child)
        if not node_queue[0].label == 'TOP' and not node_queue[0].label == 'TURN':
            counter += 1
        node_queue = node_queue[1:]
    return counter

for pp,gg in zip(preds,golds):
    pred_total += count_nodes(pp)
    gold_total += count_nodes(gg)

print(f'pred total: {pred_total}')
print(f'gold total: {gold_total}')
      
