import os
import trees

uh = 'UH'

split = 'train'

data = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/turn_pipeline/'

tree_f = os.path.join(data,f'{split}.trees')
id_f = os.path.join(data,f'{split}.ids')
lbl_f = os.path.join(data,f'{split}.lbl')

treebank, ids = trees.load_trees_with_idx(tree_f,id_f)

lbls = [[int(lbl) for lbl in line.split()] for line in open(lbl_f,'r').readlines()]


final_uhs = 0
initial_uhs = 0
all_uhs = 0

for idnum,lbl,tr in zip(ids,lbls,treebank):
    tags = [leaf.tag for leaf in tr.leaves()]
    tags[-1] = 0
    if uh in set(tags):
        uh_idxs = [index for index, element in enumerate(tags) if element == uh]
        end_idxs = [index for index, element in enumerate(lbl) if element == 1]
        start_idxs = [idx+1 for idx in end_idxs]
        final_uhs += len(set(uh_idxs).intersection(set(end_idxs)))
        initial_uhs += len(set(uh_idxs).intersection(set(start_idxs)))
        all_uhs += len(uh_idxs)
import pdb;pdb.set_trace()
        
    
