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_pipeline'
out = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/code/self_attn_speech_parser/output/turn_pipeline'

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 = True
split = 'dev'
golds = load_lines(os.path.join(data,f'{split}.lbl'))
if SPEECH:
    pred_f = f'turn_sp_seg_72240_focal_loss_no_alpha_gamma3_final_one_dev=87.61.pt_{split}_predicted.txt'
else:
    pred_f = 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,pred_f)
    
    
ids = load_lines(os.path.join(data,f'{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

    print('CHECK to make sure that setting the final gold/pred value to 0 works here')
    import pdb;pdb.set_trace()
    gold[-1] = '0'
    pred[-1] = '0'
    

    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'{split}_partition.pickle'))
    turn2pitch = load_pickle(os.path.join(data,f'{split}_pitch.pickle'))
    turn2fbank = load_pickle(os.path.join(data,f'{split}_fbank.pickle'))
    turn2pause = load_pickle(os.path.join(data,f'{split}_pause.pickle'))
    turn2dur = load_pickle(os.path.join(data,f'{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]
        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]            

            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 #######
    """ # For now only write the silver sp feats and see if that works
    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')
