import trees
import os

splits = ['train','dev','test']

datadir = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/turn_pause_dur_fixed'
outdir = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/code/self_attn_speech_parser/output/turn_pause_dur_fixed'

for split in splits:
    print(f'processing {split}...')
    tree_f = os.path.join(datadir,f'turn_{split}_medium.trees')
    id_f = os.path.join(datadir,f'turn_{split}_sent_ids_medium.txt')

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

    out_plaintext = os.path.join(datadir,'seg',f'turn_{split}.txt')
    out_lbl = os.path.join(datadir,'seg',f'turn_{split}.lbl')
    out_ids = os.path.join(datadir,'seg',f'turn_{split}.ids')    

    with open(out_plaintext,'w') as ftxt:
        with open(out_lbl,'w') as flbl:
            with open(out_ids,'w') as fid:
                for idnum,tree in zip(ids,tree_list):
                    words = [leaf.word for leaf in tree.leaves()]
                    lbls = [0 for wd in words]
                    children = tree.children
                    cumulative_len = 0
                    for child in children:
                        cumulative_len += len(list(child.leaves()))
                        if cumulative_len < len(words): lbls[cumulative_len-1] = 1
                    lbls = [str(i) for i in lbls]

                    ftxt.write(' '.join(words)+'\n')
                    flbl.write(' '.join(lbls)+'\n')
                    fid.write(idnum+'\n')
                
                
        
                           
    
