import trees
import os
import numpy as np
import pickle

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

datadir = '/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_opensmile'
outdir = '/disk/scratch1/s1899827/prosody_nlp/code/self_attn_speech_parser/output/turn_no_voicing'


for split in splits:
    print(f'processing {split}...')
    id2lbl = {}
    tree_f = f'sp_glove_dp.3_72240_{split}=86.48.pred'
    #tree_f = f'turn_sp_parse_72240_dummy_penalty.5_dev=86.36.pt_{split}.pred'

    id_f = os.path.join(datadir,f'turn_{split}.ids')


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


    #  out_pickle = os.path.join(datadir,f'{tree_f}_predlbl.pickle')
    
    tree_f = os.path.join(outdir,tree_f)
    tree_list,ids = trees.load_trees_with_idx(tree_f,id_f)

    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[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')
                    id2lbl[idnum] = np.array(lbls)
                

    with open(out_pickle,'wb') as f:
        pickle.dump(id2lbl,f)
    
                           
    
