"""
Document the distribution of expansions for each non-terminal.

I'll do this two ways. Given this example:

S -> NP VP

1) I'll count this as one occurence of an atom called 'NP VP'
2) I'll count one occurence of 'NP' and one of 'VP'

"""
import os
import trees

datadir = '/afs/inf.ed.ac.uk/group/project/prosody/prosody_nlp/data/input_features/turn_opensmile/'
tree_f = os.path.join(datadir,'turn_train.trees')
id_f = os.path.join(datadir,'turn_train.ids')

tree_list,ids = trees.load_trees_with_idx(tree_f,id_f)

exp_freq_gestalt = {} # treat each RHS of a rule as a non-compositional whole
exp_freq = {} # count each RHS as made up of atoms

for idnum,tree in zip (ids,tree_list):
    #print(tree.linearize())
    nodeq = [tree]
    while nodeq:
        curr_node = nodeq[0]
        curr_lbl = curr_node.label
        nodeq = nodeq[1:]
        curr_child = [node for node in curr_node.children if not node.is_leaf()]
        #print(curr_child)
        #import pdb;pdb.set_trace()
        if not curr_lbl in exp_freq_gestalt: exp_freq_gestalt[curr_lbl] = {}
        if not curr_lbl in exp_freq: exp_freq[curr_lbl] = {}
        for child in curr_child:
            child_lbl = child.label
            if child_lbl in exp_freq[curr_lbl]:
                exp_freq[curr_lbl][child_lbl] += 1
            else:
                exp_freq[curr_lbl][child_lbl] = 1
        child_gestalt = ' '.join([child.label for child in curr_child])
        if child_gestalt in exp_freq_gestalt[curr_lbl]:
            exp_freq_gestalt[curr_lbl][child_gestalt] += 1
        else:
            exp_freq_gestalt[curr_lbl][child_gestalt] = 1
            
        nodeq.extend(curr_child)

print(exp_freq)
import pdb;pdb.set_trace()


# TODO NEXT: plot dist for each LHS.
# (pickle the dict, scp it to local, run pyplot lib
