# 1. Load in the parses
# 2. Load in the segmentations
# 3. Eliminate all constituents that cross a misplaced boundary
#       Do this by adding fencepost annotations to the node objects and checking for overlap
# 4. OPTIONAL: add in constituents to be daughters of TURN if the alternative would be sub-optimal

import numpy as np
import os
import trees

splits = ['dev','train']

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



def all_nodes(node):
    nodes = [node]
    try:
        children = node.children
        for child in children:
            nodes.extend(all_nodes(child))
    except:
        
        return nodes
    
    return nodes
    


for split in splits:
    print(split)
    goldf = os.path.join(datadir,f'turn_{split}.trees')
    idf = os.path.join(datadir,f'turn_{split}.ids')

    predf = os.path.join(outdir,f'turn_sp_seg_72240_focal_loss_no_alpha_gamma2_dev=76.90.pt_{split}_predicted.txt')

    preds = [[int(idx) for idx in line.strip().split()] for line in open(predf,'r').readlines()]

    gold_trees,ids = trees.load_trees_with_idx(goldf,idf)
    id2tree = dict(zip(ids,gold_trees))
    tree2id = dict(zip(gold_trees,ids))
    
        

    for tree,pred in zip(gold_trees,preds):
        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


                        for idx in pred_idxs:
                            if idx < end and idx > start and not child_lbl in ['TURN','TOP']:
                                # REMOVE NODE
                                #print(f"REMOVING NODE WITH LBL {child_lbl}")
                                print(tree2id[tree])
                                remove_child = True
                        if remove_child:


                            revised_children.extend(child.children)
                        else:
                            revised_children.append(child)
                    curr_node.set_children(revised_children)
                    node_queue.extend(revised_children)
                    node_queue = node_queue[1:]


    silverf = os.path.join(datadir,f'silver_{split}.trees')
    with open(silverf,'w') as f:
        for idnum in ids:
            f.write(id2tree[idnum].linearize())
            f.write('\n')


