import trees
import os
import pandas as pd
import random
import pprint

pp = pprint.PrettyPrinter(indent=4)


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

goldf = os.path.join(datadir,'turn_dev.trees')
idf = os.path.join(datadir,'turn_dev.ids')
joint_predf = os.path.join(joint_outdir,'turn_sp_glove_72240_part_fixed_dev=86.34.pt_dev_predicted.txt')
pipe_predf = os.path.join(pipe_outdir,'sp_glove_72240_dev=76.12.pt_parsed.txt.turns')



gold_trees,ids = trees.load_trees_with_idx(goldf,idf)
pred_trees,ids = trees.load_trees_with_idx(joint_predf,idf)

"""
Collect the freq of every node that is a direct daughter of TURN
"""
goldnode2freq = {}
prednode2freq = {}

oversplit_preds = []
oversplit_golds = []

pred_only_child = []

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

def get_child_lbls(node):
    lbls = set()
    for child in node.children:
        try:
            lbl = child.label
        except:
            lbl = child.tag
        lbls.add(lbl)
    return lbls
        
for pred,gold in zip(pred_trees,gold_trees):
    assert gold.label=='TURN'
    assert pred.label=='TURN'
    gold_children = gold.children
    pred_children = pred.children
    pred_only_lbls = list(get_child_lbls(pred) - get_child_lbls(gold))
    pred_only_child.append(pred_only_lbls)
    if len(gold_children) < len(pred_children):
        oversplit_preds.append(pred)
        oversplit_golds.append(gold)
        
    for child in gold_children:
        try:
            lbl = child.label
        except:
            lbl = child.tag            
        if lbl in goldnode2freq:
            goldnode2freq[lbl] += 1
        else:
            goldnode2freq[lbl] = 1
    for child in pred_children:
        try:
            lbl = child.label
        except:
            lbl = child.tag

        if lbl in prednode2freq:
            prednode2freq[lbl] += 1
        else:
            prednode2freq[lbl] = 1


"""

Here put in a graph or other data print out of the freqs
"""
print('Frequencies for gold parses')
pp.pprint(goldnode2freq)
print('Frequencies for predicted parses')
pp.pprint(prednode2freq)

pred_only_nodes = set(prednode2freq.keys()) - set(goldnode2freq.keys())
print('nodes that are direct daughters of TURN in the predicted parses, but not gold parses')
print(pred_only_nodes)
import pdb;pdb.set_trace()



"""
Next, load the ex where num children is higher for pred than for gold
Sample 10 percent of these randomly and write to csv
load csv in spreadsheet and annotate
"""


rand_idx = list(set(range(len(oversplit_preds))))
random.shuffle(rand_idx)
rand_idx = rand_idx[:int(len(oversplit_preds)*.1)]

tmp_pred = []
tmp_gold = []
tmp_pred_lbls = []

for idx in rand_idx:

    tmp_pred.append(oversplit_preds[idx].linearize())
    tmp_gold.append(oversplit_golds[idx].linearize())
    tmp_pred_lbls.append(' '.join(pred_only_child[idx]))
    
df = pd.DataFrame(tmp_pred,columns=['pred'])
df['gold'] = tmp_gold
df['children_of_TURN_unique_to_pred'] = tmp_pred_lbls
df.to_csv('comparison.csv',sep='\t')

import pdb;pdb.set_trace()
