import os
import pickle
import numpy as np


split = 'train'

featdir = f'/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/turn_opensmile/'
pitchf = os.path.join(featdir,f'turn_{split}_pitch.pickle')
voicingf = os.path.join(featdir,f'turn_{split}_voicing.pickle')
osf = os.path.join(featdir,f'turn_{split}_opensmile.pickle')

turn2pitch = pickle.load(open(pitchf,f'rb'))
#turn2voicing = pickle.load(open(voicingf,f'rb'))
turn2os = pickle.load(open(osf,f'rb'))

new_turn2voicing = {}

for turn in turn2pitch:
    pp = turn2pitch[turn]
    #vv = turn2voicing[turn]
    oo = turn2os[turn]

    v_from_o = np.take(oo,[1,3,5],axis=0)
    new_turn2voicing[turn] = v_from_o
    print(oo.shape)
    print(pp.shape)
    print('-'*15)
    
    assert oo.shape[-1] == pp.shape[-1]

with open(voicingf,'wb') as f:
    pickle.dump(new_turn2voicing,f)


