import trees
import os
import numpy as np
import copy


def bin2spans(seg):
    """
    Convert a binary label to a list of spans representing each child node of the top node
    """
    counter = 0
    spans = []
    prev = 0
    for ii in seg:
        if ii == 0:
            pass
        else:
            spans.append((prev,counter))
            prev = counter + 1
        counter += 1
    return spans

def tree2spans(tree):
    """
    Convert a tree to a list of spans representing each child node of the top node
    """
    spans = []
    for child in tree.children:
        spans.append((child.start,child.end))
    return spans

def enforce_new_seg(treebank,segs):
    # 1: Convert segmentation to node spans
    # 2: Check the immediate child nodes to see if any covers that same span
    #    If so, pass
    #    If not, then add node, and make its children be the immediate children of TURN that are dominated by that span
    for seg,tree in zip(segs,treebank):
        tree.add_fenceposts()
        
        new_spans = bin2spans(seg)
        for nspan in new_spans:
            add_span = True
            children = tree.children 
            grandchildren = []
            for idx,child in enumerate(children):
                if child.start==nspan[0] and child.end==nspan[1]:
                    add_span = False
                elif child.start >= nspan[0] and child.end <= nspan[1]:
                    grandchildren.append(idx)
            if add_span:

                grandchildren_nodes = [tree.children[i] for i in grandchildren]
                new_node = trees.InternalTreebankNode(label='S',children=grandchildren_nodes,start=nspan[0],end=nspan[1])
                new_children = []
                span_added = False
                for j in range(len(tree.children)):
                    if not j in grandchildren:
                        new_children.append(tree.children[j])
                    elif not span_added:
                        new_children.append(new_node)
                        span_added = True
                tree.children = new_children

                    

def remove_inconsistent_seg(segs,segs_from_parse):
    """
    Gets rid of segmentations in the prediction that aren't present in the original parse
    """
    counter = 0
    for ss,pp in zip(segs,segs_from_parse):

        ss_copy = copy.deepcopy(ss)
        # Where ss==1 and pp==0, set ss=0
        # (actually only need to check pp)
        pp_mask = pp==0
        ss[pp_mask]=0
        if not np.array_equal(ss_copy,ss):
            counter += 1
    return segs

def test_remove_inconsistent_seg(segs,segs_from_parse):
    for ss,pp in zip(segs,segs_from_parse):
        pp_mask = pp==0
        try:
            assert np.all(ss[pp_mask]==0)
        except:
            print("Failed test! Segmentations aren't consistent with parse")
            import pdb;pdb.set_trace()

    
def main():
    outdir = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/code/self_attn_speech_parser/output/turn_opensmile'
    datadir = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/turn_opensmile'
    parse_f = os.path.join(outdir,'turn_nonsp_72240_dummy_penalty.5_dev=85.33.pt_dev.pred')#turn_sp_parse_72240_dummy_penalty.5_dev=86.36.pt_dev_predicted.txt')
    id_f = os.path.join(datadir,'turn_dev.ids')
    seg_f = os.path.join(outdir,'nonsp_predlbl_dev=90.35.pred')#sp_predlbl_dev=90.26.pred')
    parse_seg_f = os.path.join(datadir,'turn_nonsp_72240_dummy_penalty.5_dev=85.33.pt_dev.pred.lbl')#turn_sp_parse_72240_dummy_penalty.5_dev=86.36.pt_dev_predicted.txt.lbl.final_one')
    
    treebank,ids = trees.load_trees_with_idx(parse_f,id_f)

    segs = [np.array([int(i) for i in line.strip().split()]) for line in open(seg_f,'r').readlines()]
    segs_from_parse =  [np.array([int(i) for i in line.strip().split()]) for line in open(parse_seg_f,'r').readlines()]

    segs = remove_inconsistent_seg(segs,segs_from_parse)
    test_remove_inconsistent_seg(segs,segs_from_parse)
                  
    enforce_new_seg(treebank,segs)

    out = os.path.join(outdir,seg_f+'.reparsed')
    with open(out,'w') as f:
        for tree in treebank:
            f.write(tree.linearize())
            f.write('\n')
        
if  __name__=="__main__":
    main()
