import os
import pickle
import numpy as np
import trees
import math

data = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/turn_opensmile'
out = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/code/self_attn_speech_parser/output/silver_reparse'

def load_lines(filename):
    return [l.strip() for l in open(filename,'r').readlines()]

def load_pickle(filename):
    with open(filename,'rb') as f:
        return pickle.load(f)

SPEECH = False
split = 'dev'
golds = load_lines(os.path.join(data,f'turn_{split}.lbl'))
if SPEECH:
    #preds = load_lines(os.path.join(out,'turn_sp_seg_72240_dev=99.71.pt_turn_dev_predicted.txt'))
    #preds = load_lines(os.path.join(out,'turn_sp_seg_72240_dev=76.12.pt_turn_dev_predicted.txt'))
    preds = load_lines(os.path.join(out,f'turn_nonsp_seg_72240_focal_loss_no_alpha_gamma3_final_one_dev=87.61.pt_{split}_predicted.txt'))
    #preds = load_lines(os.path.join(out,'turn_sp_seg_72240_focal_loss_no_alpha_gamma3_final_one_dev=89.84.pt_train_predicted.txt'))
else:
    preds = load_lines(os.path.join(out,f'turn_nonsp_seg_72240_focal_loss_no_alpha_gamma3_final_one_dev=87.61.pt_{split}_predicted.txt'))

    
ids = load_lines(os.path.join(data,f'turn_{split}.ids'))
turn2sent = load_pickle(os.path.join(data,'..','turn2sent.pickle'))

############## Make a dict of turn id to list of tokens.

tree_file = os.path.join(data,'..','sentence_pause_dur_fixed',f'{split}.trees')
id_file =  os.path.join(data,'..','sentence_pause_dur_fixed',f'{split}_sent_ids.txt')

treebank, all_sents = trees.load_trees_with_idx(tree_file,id_file) # EKN trees get loaded in as trees here
id2senttree = dict(zip(all_sents,treebank))


def get_toks_from_tree(tree):
    return [leaf.word for leaf in tree.leaves()]

def get_tags_from_tree(tree):
    return [leaf.tag for leaf in tree.leaves()]

turn2toks = {}
turn2tags = {}
for turn in turn2sent:
    if len(turn2sent[turn])==1:
        sent = turn2sent[turn][0]
        if sent in id2senttree:
            turn2toks[turn] = get_toks_from_tree(id2senttree[sent])
            turn2tags[turn] = get_tags_from_tree(id2senttree[sent])
    else:
        sents = turn2sent[turn]
        for sent in sents:
            if sent in id2senttree:
                if not turn in turn2toks:
                    turn2toks[turn] = get_toks_from_tree(id2senttree[sent])
                    turn2tags[turn] = get_tags_from_tree(id2senttree[sent])
                else:
                    turn2toks[turn].extend(get_toks_from_tree(id2senttree[sent]))
                    turn2tags[turn].extend(get_tags_from_tree(id2senttree[sent]))
        
### STEP 1 #################################################
# correctly segmented turns -> add sent IDS to list
# incorrectly segmented turns 
#    -> combine sent ids and add to list
#    -> split sent ids and add to list
############################################################

total_golds = len(golds)
for gold in golds:
    total_golds += sum([int(i) for i in gold.split()])

print(total_golds)
    


usable_sent_ids = []

corrects = 0
unders = 0
overs = 0
gold_unders = 0
pred_sents = {} # stores the incorrecty segmented predicted sents
pred_sent_tags = {}

for idnum,gold,pred in zip(ids,golds,preds):
    g_list = np.array([int(i) for i in gold.split()])
    total_golds += np.sum(g_list) + 1
    if gold==pred:
        p_list = np.array([int(i) for i in pred.split()])
        corrects += p_list.sum() + 1
        usable_sent_ids.extend(turn2sent[idnum])
    else:
        turn_toks = turn2toks[idnum]
        turn_tags = turn2tags[idnum]
        g_list = np.array([int(i) for i in gold.split()])
        p_list = np.array([int(i) for i in pred.split()])
        wd_idx = 0
        cumulative_toks = []
        cumulative_tags = []
        pred_sent_start = 0
        for gg,pp in zip(g_list,p_list):
            if pp == 0:
                cumulative_toks.append(turn_toks[wd_idx])
                cumulative_tags.append(turn_tags[wd_idx])
            if pp == 1:
                cumulative_toks.append(turn_toks[wd_idx])
                cumulative_tags.append(turn_tags[wd_idx])
                pred_sent_id = f'{idnum}+{pred_sent_start}+{wd_idx}'
                pred_sents[pred_sent_id] = cumulative_toks
                pred_sent_tags[pred_sent_id] = cumulative_tags
                pred_sent_start = wd_idx+1
                cumulative_toks = []
                cumulative_tags = []
            if wd_idx==len(turn_toks)-1 and pp==0:
                pred_sent_id = f'{idnum}+{pred_sent_start}+{wd_idx}'
                pred_sents[pred_sent_id] = cumulative_toks
                pred_sent_tags[pred_sent_id] = cumulative_tags
                pred_sent_start = wd_idx+1
                cumulative_toks = []
                cumulative_tags = []
            wd_idx += 1


usable_sent_ids = set(usable_sent_ids)

##### STEP 2 ##############################
# Put sent ids and words into new files
###########################################


outid2tree = {}
for idnum in id2senttree:
    if idnum in usable_sent_ids:
        outid2tree[idnum] = id2senttree[idnum].linearize()
        usable_sent_ids.remove(idnum)

for pred in pred_sents:
    curr_wds = pred_sents[pred]
    curr_tags = pred_sent_tags[pred]
    outid2tree[pred] = '(TURN '+' '.join([f'({tag} {wd})' for tag,wd in zip(curr_tags,curr_wds)])+')'

   
"""
for idnum in usable_sent_ids:
    if ':' in idnum:
        curr_sents = idnum.split(':')
        trees = [id2senttree[sent].linearize() for sent in curr_sents]
        new_tree = '(TURN '+' '.join(trees)+')'
        outid2tree[idnum] = new_tree
    if '+' in idnum:
        original_id,start,end = idnum.split('+')
        start = int(start)
        end = int(end)
        wds = [leaf.word for leaf in id2senttree[original_id].leaves()]
        tags = [leaf.tag for leaf in id2senttree[original_id].leaves()]
        curr_wds = wds[start:end]
        curr_tags = tags[start:end]
        new_tree = '(TURN '+' '.join([f'({tag} {wd})' for tag,wd in zip(curr_tags,curr_wds)])+')'
        # Note: this is a dummy tree just to get the parser to parse the sentence;
        # we'll use the real turn for eval.
        outid2tree[idnum] = new_tree
"""

##### STEP 3 ##############################
# Remake acoustic features
###########################################
outfeats = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/silver'

if SPEECH:
    feats = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/sentence_pause_dur_fixed'

    kaldi_feats = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/kaldi_feats/'
    sent2pause = pickle.load(open(os.path.join(feats,f'{split}_pause.pickle'),'rb'))
    sent2part = pickle.load(open(os.path.join(feats,f'{split}_partition.pickle'),'rb'))
    sent2dur = pickle.load(open(os.path.join(feats,f'{split}_duration.pickle'),'rb'))
    sent2fbank = pickle.load(open(os.path.join(feats,f'{split}_fbank.pickle'),'rb'))
    sent2pitch = pickle.load(open(os.path.join(feats,f'{split}_pitch.pickle'),'rb'))
    sent2term = pickle.load(open(os.path.join(feats,'..','sent2term.pickle'),'rb'))
    term2feats = pickle.load(open(os.path.join(feats,'..',f'term2feats.pickle'),'rb'))

    turn2part = load_pickle(os.path.join(data,f'turn_{split}_partition.pickle'))
    turn2pitch = load_pickle(os.path.join(data,f'turn_{split}_pitch.pickle'))
    turn2fbank = load_pickle(os.path.join(data,f'turn_{split}_fbank.pickle'))
    turn2pause = load_pickle(os.path.join(data,f'turn_{split}_pause.pickle'))
    turn2dur = load_pickle(os.path.join(data,f'turn_{split}_duration.pickle'))

    
    newsent2pause = {}
    newsent2dur = {}
    newsent2pitch = {}
    newsent2fbank = {}
    newsent2part = {}

    def times2frames(stime,etime,audio_len):
        sframe = max(0,math.floor(stime*100))
        eframe = min(math.ceil(etime*100),audio_len)
        return sframe,eframe


    for idnum in outid2tree:
        if idnum in id2senttree:
            newsent2pause[idnum] = sent2pause[idnum]
            newsent2part[idnum] = sent2part[idnum]
            newsent2pitch[idnum] = sent2pitch[idnum]
            newsent2dur[idnum] = sent2dur[idnum]
            newsent2fbank[idnum] = sent2fbank[idnum]

            """
            elif ':' in idnum:
            sents = idnum.split(':')
        
            durs = [sent2dur[sent] for sent in sents]
            newsent2dur[idnum] = np.concatenate(durs,axis=1)

            parts = []
            pauses = {'pause_bef':[],
                      'pause_aft':[]}

            for sent in sents:
                parts.extend(sent2part[sent])
                pauses['pause_bef'].extend(sent2pause[sent]['pause_bef'])
                pauses['pause_aft'].extend(sent2pause[sent]['pause_aft'])
            
            newsent2pause[idnum] = pauses
            newsent2part[idnum] = parts

            ssent = sents[0]
            esent = sents[-1]
            # Pitch:

            conv,spk,_ = sent.split('_')

            pitch_dict = pickle.load(open(os.path.join(kaldi_feats,'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)

            curr_pitch = pitch_data[:,sframe:eframe]
            newsent2pitch[idnum] = curr_pitch
        
            fbank_data = pickle.load(open(os.path.join(kaldi_feats,'swbd_fbank_energy','normed_fbanks',f'{conv}-{spk}.pickle'),'rb'))
            curr_fbank = fbank_data[:,sframe:eframe]
            newsent2fbank[idnum] = curr_fbank

            assert curr_fbank.shape[-1]==curr_pitch.shape[-1]
            assert len(parts) == len(pauses['pause_bef'])
            assert len(parts) == len(pauses['pause_aft'])
            assert len(parts) == newsent2dur[idnum].shape[-1]

            """
            #elif '+' in idnum:
        else:
            turn,start,end = idnum.split('+')
            conv,spk,_ = turn.split('_')

            start = int(start)
            end = int(end)
            
            partition = turn2part[turn]
            pitch = turn2pitch[turn]
            fbank = turn2fbank[turn]
            dur = turn2dur[turn]
            pause = turn2pause[turn]

            start_frame = partition[start][0]
            end_frame = partition[end][-1]
            newsent2pitch[idnum] = pitch[:,start_frame:end_frame]
            newsent2fbank[idnum] = fbank[:,start_frame:end_frame]
            newsent2part[idnum] = [(part[0]-start_frame,part[1]-start_frame) for part in partition[start:end+1]]
            newsent2dur[idnum] = dur[:,start:end+1]
            newsent2pause[idnum] = {}
            newsent2pause[idnum]['pause_bef'] = pause['pause_bef'][start:end+1]
            newsent2pause[idnum]['pause_aft'] = pause['pause_aft'][start:end+1]            

            """
            terms = sent2term[sent]
            orths = [term2feats[term]['orth'] for term in terms][start:end]

            sterm = terms[start]
            eterm = terms[end]
            sorth = term2feats[sterm]['orth']
            eorth = term2feats[eterm]['orth']

            stime = term2feats[sterm]['start']
            etime = term2feats[eterm]['end']        

            pitch_dict = pickle.load(open(os.path.join(kaldi_feats,'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
            sframe,eframe = times2frames(stime,etime,audio_len)
            eframe = eframe
            
            curr_pitch = pitch_data[:,sframe:eframe]
            newsent2pitch[idnum] = curr_pitch
        
            fbank_data = pickle.load(open(os.path.join(kaldi_feats,'swbd_fbank_energy','normed_fbanks',f'{conv}-{spk}.pickle'),'rb'))
            curr_fbank = fbank_data[:,sframe:eframe]
            newsent2fbank[idnum] = curr_fbank

            
            pause_bef = sent2pause[sent]['pause_bef'][start:end+1]
            pause_aft = sent2pause[sent]['pause_aft'][start:end+1]
            newsent2pause[idnum] = {'pause_bef':pause_bef,
                                    'pause_aft':pause_aft}
            """
            try:
                assert len(newsent2part[idnum]) == len(pred_sents[idnum])
            except:
                import pdb;pdb.set_trace()
            assert len(newsent2part[idnum]) == len(newsent2pause[idnum]['pause_bef'])
            assert len(newsent2part[idnum]) == len(newsent2pause[idnum]['pause_aft'])
            assert len(newsent2part[idnum]) == newsent2dur[idnum].shape[-1]

            
    assert len(newsent2part) == len(newsent2fbank)
    assert len(newsent2part) == len(newsent2pitch)
    assert len(newsent2part) == len(newsent2pause)
    assert len(newsent2part) == len(newsent2dur)
    assert len(newsent2part) == len(outid2tree)
    print(f'num pred sents: {len(newsent2part)}')
    ## STEP 4: write all the files #######
    with open(os.path.join(outfeats,f'{split}_fbank.pickle'),f'wb') as f:
        pickle.dump(newsent2fbank,f)
    with open(os.path.join(outfeats,f'{split}_pitch.pickle'),f'wb') as f:
        pickle.dump(newsent2pitch,f)
    with open(os.path.join(outfeats,f'{split}_pause.pickle'),f'wb') as f:
        pickle.dump(newsent2pause,f)
    with open(os.path.join(outfeats,f'{split}_partition.pickle'),f'wb') as f:
        pickle.dump(newsent2part,f)
    with open(os.path.join(outfeats,f'{split}_duration.pickle'),f'wb') as f:
        pickle.dump(newsent2dur,f)


if SPEECH:
    trees_outfile = f'{split}_for_parser.trees.nonsp'
    ids_outfile = f'{split}_sent_ids_for_parser.txt.nonsp'
else:
    trees_outfile = f'{split}_nonsp_for_parser.trees.nonsp'
    ids_outfile = f'{split}_nonsp_sent_ids_for_parser.txt.nonsp'


with open(os.path.join(outfeats,trees_outfile),f'w') as ftree:
    with open(os.path.join(outfeats,ids_outfile),f'w') as fid:
        for idnum in outid2tree:
            tree = outid2tree[idnum]
            if tree==f'(TURN )':
                import pdb;pdb.set_trace()
            ftree.write(tree)
            ftree.write(f'\n')
            fid.write(idnum)
            fid.write(f'\n')
