import pickle
import os
from prep_input_dicts import times2frames,norm_energy_by_side,make_sent2part_and_times,load_orth2meandur_json
import numpy as np
import pandas as pd

# 0. Set paths
data_dir = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/'
sent_dir = os.path.join(data_dir,'sentence_short')
#sent_dir = os.path.join(data_dir,'sentence_pause_dur_fixed')
turn_dir = os.path.join(data_dir,'turn_part_fixed')
kaldi_feat_dir = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/kaldi_feats/'
fbank_dir = os.path.join(kaldi_feat_dir,'swbd_fbank_energy')
pitchpov_dir = os.path.join(kaldi_feat_dir,'swbd_pitch_pov')
opensmile_dir = os.path.join(data_dir,'opensmile')

split = 'train'
#split = 'dev'
#split = 'test' 

# 1. Load intermediate dict files that are potentially helpful
print('Loading helper dicts ...')
sent2pw = pickle.load(open(os.path.join(data_dir,f'sent2pw.pickle'),'rb'))
sent2term = pickle.load(open(os.path.join(data_dir,f'sent2term.pickle'),'rb'))
term2pw = pickle.load(open(os.path.join(data_dir,f'term2pw.pickle'),'rb'))
turn2sent = pickle.load(open(os.path.join(data_dir,f'turn2sent.pickle'),'rb'))
sent2turn = pickle.load(open(os.path.join(data_dir,f'sent2turn.pickle'),'rb'))
term2feats = pickle.load(open(os.path.join(data_dir,f'term2feats.pickle'),'rb'))
#pw2feats = pickle.load(open(os.path.join(data_dir,f'pw2feats.pickle'),'rb'))

phone2meandur = pickle.load(open(os.path.join(data_dir,'phone2meandur.pickle'),'rb'))



# 2. Files for the parser that are split into train/dev/test
print('Loading sentence-based dicts ...')
sent2part = pickle.load(open(os.path.join(sent_dir,f'{split}_partition.pickle'),'rb'))
sent2dur = pickle.load(open(os.path.join(sent_dir,f'{split}_duration.pickle'),'rb'))
sent2rawdur = pickle.load(open(os.path.join(sent_dir,f'{split}_raw_duration.pickle'),'rb'))
sent2fbankenergy = pickle.load(open(os.path.join(sent_dir,f'{split}_fbank.pickle'),'rb'))
sent2pitchpov= pickle.load(open(os.path.join(sent_dir,f'{split}_pitch.pickle'),'rb'))
#sent2opensmile=pickle.load(open(os.path.join(sent_dir,f'{split}_opensmile.pickle'),'rb'))
sent2opensmile=pickle.load(open(os.path.join(sent_dir,f'{split}_voicing.pickle'),'rb'))
sent2pause = pickle.load(open(os.path.join(sent_dir,f'{split}_pause.pickle'),'rb'))
with open(os.path.join(sent_dir,f'{split}_sent_ids.txt'),'r') as f:
    sent_ids = [l.strip() for l in f.readlines()]
with open(os.path.join(sent_dir,f'{split}.trees'),'r') as f:
    trees = [l.strip() for l in f.readlines()]

sent2tree = dict(zip(sent_ids,trees))

# 3. Create dicts that are turn based:
turn2part = {}
turn2dur = {}
turn2fbank = {}
turn2pitch = {}
turn2pause = {}
turn2tree = {}
turn2opensmile = {}


unsorted_sents = sent_ids
sent_ids = sorted(list(sent2part.keys()))
sents_set = set(sent_ids)

sents_in_turns = []



print('Sorting turns ...')
# 4. Separate turns where turn=sent from others
def sort_turns(turn2sent,sent2part):
    multisent_turns = []
    singlesent_turns = []
    single_sents = []
    for turn in turn2sent:
        sents = turn2sent[turn]
        sents_exist = True
        for sent in sents:
            if not sent in sent2part:
                sents_exist = False
        if sents_exist:
            sents_in_turns.extend(sents)
            if len(turn2sent[turn]) > 1:
                multisent_turns.append(turn)
            elif len(turn2sent[turn]) == 1:
                singlesent_turns.append(turn)
                single_sents.extend(turn2sent[turn])
    multisent_turns = sorted(multisent_turns)
    singlesent_turns = sorted(singlesent_turns)
    single_sents = sorted(single_sents)
    return multisent_turns,singlesent_turns,single_sents

multisent_turns,singlesent_turns,single_sents = sort_turns(turn2sent,sent2part)


#4b.  Create turn2part from scratch
conv2sent = {}
for sent in sent2turn:
    conv,_,_ = sent.split('_')
    if conv in conv2sent:
        conv2sent[conv].append(sent)
    else:
        conv2sent[conv] = [sent]

convs = list(conv2sent.keys())
alignment_dir = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/swbd_word_times'

def load_pw2feats(convs,alignment_dir,orth2meandur):
    # Load a dict from phonword to features for a given conversation (both speakers)
    spks = ['A','B']
    pw2feats = {}
    for conv in convs:
        pw2feats[conv]={}
        for spk in spks:
            # Load pause and norm feats that are generated by the defunct 'tree alignment' code
            align_file = f'word_times_{conv}{spk}.pickle'
            with open(os.path.join(alignment_dir,align_file),'rb') as f:
                align_dict = pickle.load(f)
            for wd in align_dict:
                pw2feats[conv][wd] = {}
                pw2feats[conv][wd]['phones'] = align_dict[wd]['phones']
                pw2feats[conv][wd]['start_time'] = align_dict[wd]['start_time']
                pw2feats[conv][wd]['end_time'] = align_dict[wd]['end_time']
                pw2feats[conv][wd]['phone_end_times'] = align_dict[wd]['phone_end_times']
                pw2feats[conv][wd]['phone_start_times'] = align_dict[wd]['phone_start_times']
                #pw2feats[wd]['mean_dur'] = orth2meandur[pw2feats[pw_id]['text']]

    return pw2feats

avg_dur_dir = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data'
wordstats_json = 'avg_word_stats.json'
orth2meandur,orth2freq = load_orth2meandur_json(avg_dur_dir,wordstats_json)



pw2feats = load_pw2feats(convs,alignment_dir,orth2meandur)



turn2times = {}
turn2term = {}
for turn in turn2sent:
    sents = turn2sent[turn]
    if len(sents) == 1:
        if sents[0] in sent2term and sents[0] in sents_in_turns:
            turn2term[turn] = sent2term[sents[0]]
    else:
        terms = []
        for sent in sents:
            if sent in sent2term  and sent in sents_in_turns:
                terms.extend(sent2term[sent])

        if terms:
            turn2term[turn] = terms

conv2turn = {}
for turn in turn2term:
    conv,_,_ = turn.split('_')
    if conv in conv2turn:
        conv2turn[conv].append(turn)
    else:
        conv2turn[conv] = [turn]


for conv in conv2turn:
    curr_turns = conv2turn[conv]
    conv_pw2feats = pw2feats[conv]
    turn2part,turn2times = make_sent2part_and_times(curr_turns,turn2part,turn2times,turn2term,term2pw,conv_pw2feats,orth2meandur,term2feats,phone2meandur)


"""
# 4b. Create files of sent ids/trees that is only the sents covered by turns.
sents_in_turns = set(sents_in_turns)
short_sent_dir = os.path.join(sent_dir,'..','sentence_short')
short_sent_ids = open(os.path.join(short_sent_dir,f'{split}_sent_ids.txt'),'w')
short_sent_trees = open(os.path.join(short_sent_dir,f'{split}.trees'),'w')
for idnum in unsorted_sents:
    tree = sent2tree[idnum]
    if idnum in sents_in_turns:
        short_sent_ids.write(idnum)
        short_sent_ids.write('\n')
        short_sent_trees.write(tree)
        short_sent_trees.write('\n')
short_sent_ids.close()
short_sent_trees.close()
"""                           

""" # Debugging residue -- figuring out which sentences were left out of the turn-based set.
print(len(sents_in_turns))
print(len(sents_set))
sent_diffs = sents_set - sents_in_turns

conv2sent = {}
for sent in sent_diffs:
    conv,_,_ = sent.split('_')
    if conv in conv2sent:
        conv2sent[conv].append(sent)
    else:
        conv2sent[conv] = [sent]
        

for conv in conv2sent:
    print(conv)
    print(conv2sent[conv])
import pdb;pdb.set_trace()
"""

redo_convs = set()

# COMMENTED OUT FOR DEBUGGING ONLY
# 5. For turns where turn=sent, just copy the sent data to a turn-based dict.
print('Copying over single-sentence turns ...')
for sent in single_sents:
    turn = sent2turn[sent]
    #turn2part[turn] = sent2part[sent]
    turn2dur[turn] = sent2dur[sent]
    turn2pause[turn] = sent2pause[sent]
    turn2pitch[turn] = sent2pitchpov[sent]
    turn2fbank[turn] = sent2fbankenergy[sent]
    turn2opensmile[turn] = sent2opensmile[sent]
    if turn2fbank[turn].shape[-1] != turn2opensmile[turn].shape[-1]:
        redo_convs.add(turn.split()[0])
    #turn2tree[turn] = f'(TURN {sent2tree[sent] })'
    tr = sent2tree[sent].replace('(TOP','')[:-1]
    turn2tree[turn] = '(TOP (TURN '+tr+'))'


# 6. Where there are 2+ sents per turn, create new feats for turn
print('Handling multi-sentence turns ...')
for turn in multisent_turns:

    sents = turn2sent[turn]

    conv,spk,_ = sents[0].split('_')
    
    # 6a. Create features you can concatenate from sent features
    pauses = {'pause_bef':[],
              'pause_aft':[]}
    #parts = []
    durs = []
    rawdurs = []
    curr_trees = ['(TOP (TURN']
    for sent in sents:
        pauses['pause_bef'].extend(sent2pause[sent]['pause_bef'])
        pauses['pause_aft'].extend(sent2pause[sent]['pause_aft'])
        #parts.extend(sent2part[sent])
        tree_no_top = sent2tree[sent].replace('(TOP','')[:-1]
        curr_trees.append(tree_no_top)
        durs.append(sent2dur[sent])
    curr_trees.append(') )')
    turn2pause[turn] = pauses
    #turn2part[turn] = parts
    turn2tree[turn] = ' '.join(curr_trees)

    # Duration:
    # recompute norm by max in turn so that it's normed by max in turn, not utterance
    normed_by_mean = np.concatenate([dur[0:1] for dur in durs],axis=1) # these are the same -- copy them over.
    raw_durs = np.concatenate([np.array(sent2rawdur[sent]) for sent in sents])
    normed_by_max = np.expand_dims(raw_durs/np.max(raw_durs),axis=0)

    

    turn2dur[turn] = np.concatenate([normed_by_mean,normed_by_max],axis=0)
    

    # 6b. Create features you have to recalculate
    ssent = sents[0]
    esent = sents[-1]
    # Pitch:
    pitch_dict = pickle.load(open(os.path.join(kaldi_feat_dir,'swbd_pitch_pov',f'{conv}-{spk}.pickle'),'rb'))
    pitch_data = np.transpose(np.array(pitch_dict[list(pitch_dict.keys())[0]]))

    audio_len = pitch_data.shape[-1] - 1
    stime = term2feats[sent2term[ssent][0]]['start']
    etime = term2feats[sent2term[esent][-1]]['end']
    sframe,eframe = times2frames(stime,etime,audio_len)
    turn_pitch = pitch_data[:,sframe:eframe]
    turn2pitch[turn] = turn_pitch
    fbank_data = pickle.load(open(os.path.join(kaldi_feat_dir,'swbd_fbank_energy','normed_fbanks',f'{conv}-{spk}.pickle'),'rb'))
    turn_fbank = fbank_data[:,sframe:eframe]
    turn2fbank[turn] = turn_fbank


    keep_feats = ['voicingFinalUnclipped_sma',
                  'logHNR_sma',
                  'pcm_zcr_sma']
    conv = conv.replace('sw','sw0')
    csv = os.path.join(opensmile_dir,f'{conv}.is13.csv')
    df = pd.read_csv(csv,sep=';')
    df = df[keep_feats]
    opensmile_feats = np.array(df).transpose()
    turn_opensmile = opensmile_feats[:,sframe:eframe]
    turn2opensmile[turn] = turn_opensmile
    if turn2fbank[turn].shape[-1] != turn2opensmile[turn].shape[-1]:
        redo_convs.add(turn.split()[0])


# 7. Save all these to pickle files
print('Pickling ...')
#pickle.dump(turn2dur,open(os.path.join(turn_dir,f'turn_{split}_duration.pickle'),'wb'))
#pickle.dump(turn2pitch,open(os.path.join(turn_dir,f'turn_{split}_pitch.pickle'),'wb'))
#pickle.dump(turn2pause,open(os.path.join(turn_dir,f'turn_{split}_pause.pickle'),'wb'))
#pickle.dump(turn2part,open(os.path.join(turn_dir,f'turn_{split}_partition.pickle'),'wb'))
#pickle.dump(turn2fbank,open(os.path.join(turn_dir,f'turn_{split}_fbank.pickle'),'wb'))
pickle.dump(turn2opensmile,open(os.path.join(turn_dir,f'turn_{split}_voicing.pickle'),'wb'))

with open(os.path.join(turn_dir,f'turn_{split}_sent_ids.txt'),'w') as f_id:
    with open(os.path.join(turn_dir,f'turn_{split}.trees'),'w') as f_tree:
        for turn in turn2tree:
            f_id.write(f'{turn}\n')
            f_tree.write(f'{turn2tree[turn]}\n')

"""
# 8. Write subsets of turn ids and trees as well:
print(f'Making turn subsets of {split} ...')
with open(os.path.join(turn_dir,f'multisent_turns_{split}_sent_ids.txt'),'w') as f_id:
    with open(os.path.join(turn_dir,f'multisent_turns_{split}.trees'),'w') as f_tree:
        for turn in multisent_turns:
            f_id.write(f'{turn}\n')
            f_tree.write(f'{turn2tree[turn]}\n')

with open(os.path.join(turn_dir,f'singlesent_turns_{split}_sent_ids.txt'),'w') as f_id:
    with open(os.path.join(turn_dir,f'singlesent_turns_{split}.trees'),'w') as f_tree:
        for turn in singlesent_turns:
            f_id.write(f'{turn}\n')
            f_tree.write(f'{turn2tree[turn]}\n')


# 9. Write sent subsets of sent ids and trees:
print(f'Making sent subsets of {split} ...')
with open(os.path.join(sent_dir,f'multisent_{split}_sent_ids.txt'),'w') as f_id_multi:
    with open(os.path.join(sent_dir,f'multisent_{split}.trees'),'w') as f_tree_multi:
        with open(os.path.join(sent_dir,f'singlesent_{split}_sent_ids.txt'),'w') as f_id_sing:
            with open(os.path.join(sent_dir,f'singlesent_{split}.trees'),'w') as f_tree_sing:
                for sent in sent_ids:
                    if sent in single_sents:
                        f_id_sing.write(sent+'\n')
                        f_tree_sing.write(sent2tree[sent]+'\n')
                    else:
                        f_id_multi.write(sent+'\n')
                        f_tree_multi.write(sent2tree[sent]+'\n')
                        
            
"""
print('Done.')

