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

TURN = ''

data = '/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_pipeline'
out = '/disk/scratch1/s1899827/prosody_nlp/code/self_attn_speech_parser/output/turn_pipeline_no_voicing_redux'

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
parser = argparse.ArgumentParser()



parser.add_argument('--predpath',
                    type=str,
                    default='sp_seg_dev=89.19.pred')
args = parser.parse_args()
predpath = args.predpath

split = predpath.split('.')[-2]
#split = 'dev'
    
preds = load_lines(os.path.join(out,predpath))
prefix = predpath.replace(f'.{split}.pred','')+'.silver.'
#prefix = predpath.replace('.pred','')+'.silver.'

golds = load_lines(os.path.join(data,f'{TURN}{split}.lbl'))

   
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')

turn_tree_file = os.path.join(data,f'{TURN}{split}.trees')
turn_id_file = os.path.join(data,f'{TURN}{split}.ids')
                       
sent_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,sent_treebank))

turn_treebank, all_turns = trees.load_trees_with_idx(turn_tree_file,turn_id_file) 
id2turntree = dict(zip(all_turns,turn_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
############################################################

def subseg(turn,idxs):
    '''
    takes in a tree and removes all nodes that are outside of the word idxs given
    
    turn: tree object (defined in trees.py)
    idxs: list of int

    makes a deep copy of turn and doesn't modify it in-place

    '''
    max_idx = max(idxs)
    min_idx = min(idxs)
    seg = copy.deepcopy(turn)
    seg.add_fenceposts()
    node_queue = [seg]
    while node_queue:
        curr_node = node_queue[0]
        if not curr_node.is_leaf():
            children = curr_node.children
            keep_children = []
            for child in children:
                if not child.is_leaf():
                    start = child.start
                    end = child.end
                    if start > max_idx or end < min_idx:
                        pass
                    else:
                        keep_children.append(child)
                else:
                    idx = child.idx
                    if idx < min_idx or idx > max_idx:
                        pass
                    else:
                        keep_children.append(child)

            node_queue.extend(keep_children)
            try:
                curr_node.set_children(keep_children)
            except:
                import pdb;pdb.set_trace()

        node_queue = node_queue[1:]
        
    if len(seg.children)==1 and not seg.children[0].is_leaf():
        seg = seg.children[0]
    else:
        seg.change_label('DUMMY')
    return seg





def remove_impossible_const(tree,pred,idnum=''):
    '''
    Removes all constituents from a tree that would be crossed by the sentence boundary predicted in pred

    modifies tree in-place, as well as returning it

    tree: tree object (defined in trees.py)
    pred: np binary array, where 1 indicates sentence boundary
    '''
    pred_idxs = np.nonzero(pred)[0].tolist()
    if sum(pred) > 0:
        tree.add_fenceposts()
        node_queue = [tree]
        while node_queue:
            curr_node = node_queue[0]
            try:
                children = curr_node.children
            except:
                children = None
                node_queue = node_queue[1:]
            if children:
                revised_children = []
                for child in children:
                    start = child.start
                    end = child.end
                    remove_child = False
                    try:
                        child_lbl = child.label
                    except:
                        child_lbl = child.tag
                    if child_lbl == 'TURN':
                        remove_child = True
                    else:
                        for idx in pred_idxs:
                            if (idx < end and idx >= start) and not child_lbl in ['TOP']:
                                # REMOVE NODE
                                remove_child = True
                    if remove_child:
                        if not child.is_leaf():
                            grandchildren = list(child.children)
                            while grandchildren:
                                curr_grandchild = grandchildren[0]
                                keep_grandchild = True
                                grandchildren = grandchildren[1:]
                                # check if grandchild is allowable. If not, add its children to the queue; if so, add it to the revised children list
                                for idx in pred_idxs:

                                    if curr_grandchild.start <= idx and curr_grandchild.end > idx:
                                        keep_grandchild = False

                                if keep_grandchild:
                                    revised_children.append(curr_grandchild)

                                elif not curr_grandchild.is_leaf():
                                    #Grandchildren.extend(curr_grandchild.children) # THIS is also a problem.
                                    grandchildren = list(curr_grandchild.children) + grandchildren
                                else:
                                    # Base case
                                    pass
                                    



                    else:
                        revised_children.append(child)
                try:
                    ''' # This was for debugging the thing where the nodes got scrambled bc of non-in-order traversal
                    if idnum=='sw4565_A_t3':
                        revised_children_print = [child.linearize()+' |' for child in revised_children]
                        print(revised_children_print)
                        import pdb;pdb.set_trace()
                    '''
                    curr_node.set_children(revised_children)
                except:
                    print('tried to add a node with no children')
                    import pdb;pdb.set_trace()
                node_queue.extend(revised_children)
                #if remove_child:
                #    import pdb;pdb.set_trace()
                node_queue = node_queue[1:]

    return tree

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

    
def tree2str(tree):
    leaves = [leaf.word for leaf in tree.leaves()]
    return ' '.join(leaves)

usable_sent_ids = []

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

silver_sents = {} # Stores the best-possible sent parse, given the predicted sent segmentation
                       
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
    gold_parse = id2turntree[idnum]
    gold = gold[:-1] +'0'
    pred = 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])
        sent_ids = turn2sent[idnum]
        for sent_id in sent_ids:
            silver_sents[sent_id] = id2senttree[sent_id]
    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()])
        gold_str = tree2str(gold_parse)
        silver_turn = remove_impossible_const(gold_parse,p_list,idnum)
        try:
            assert tree2str(silver_turn)==gold_str
        except:
            import pdb;pdb.set_trace()
        wd_idx = 0
        cumulative_toks = []
        cumulative_tags = []
        cumulative_idxs = []
        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])
                cumulative_idxs.append(wd_idx)
            if pp == 1:
                cumulative_toks.append(turn_toks[wd_idx])
                cumulative_tags.append(turn_tags[wd_idx])
                cumulative_idxs.append(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
                prev_sent_start = pred_sent_start
                pred_sent_start = wd_idx+1
                silver_sent = subseg(silver_turn,cumulative_idxs)
                cumulative_toks = []
                cumulative_tags = []
                cumulative_idxs = []
                silver_sents[pred_sent_id] = silver_sent
            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
                silver_sent = subseg(silver_turn,cumulative_idxs)

                cumulative_toks = []
                cumulative_tags = []
                cumulative_idxs = []

                silver_sents[pred_sent_id] = silver_sent
            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)])+')'

##### STEP 3 ##############################
# Remake acoustic features
###########################################

if SPEECH:
    feats = '/disk/scratch1/s1899827/prosody_nlp/data/input_features/sentence_pause_dur_fixed'
    outfeats = '/disk/scratch1/s1899827/prosody_nlp/data/input_features/turn_pipeline_DUMMY'
    kaldi_feats = '/disk/scratch1/s1899827/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'))
    #sent2voicing = pickle.load(open(os.path.join(feats,'..','sentence_short',f'{split}_voicing.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'))
    #turn2voicing = load_pickle(os.path.join(data,f'{TURN}{split}_voicing.pickle'))

    newsent2voicing = {}
    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]
            #newsent2voicing[idnum] = sent2voicing[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]
            #voicing = turn2voicing[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]
            #newsent2voicing[idnum] = voicing[:,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(newsent2voicing)
    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'{prefix}{split}_fbank.pickle'),'wb') as f:
        pickle.dump(newsent2fbank,f)
    #with open(os.path.join(outfeats,f'{prefix}{split}_voicing.pickle'),'wb') as f:
    #    pickle.dump(newsent2voicing,f)
    with open(os.path.join(outfeats,f'{prefix}{split}_pitch.pickle'),'wb') as f:
        pickle.dump(newsent2pitch,f)
    with open(os.path.join(outfeats,f'{prefix}{split}_duration.pickle'),'wb') as f:
        pickle.dump(newsent2dur,f)
    with open(os.path.join(outfeats,f'{prefix}{split}_pause.pickle'),'wb') as f:
        pickle.dump(newsent2pause,f)
    with open(os.path.join(outfeats,f'{prefix}{split}_partition.pickle'),'wb') as f:
        pickle.dump(newsent2part,f)

trees_outfile = f'{prefix}{split}.trees'
ids_outfile = f'{prefix}{split}.ids'

print(outfeats)

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

                    
