""" Compute heads of mentions. """

import logging
import re

from cort.core import spans

logger = logging.getLogger(__name__)

__author__ = 'smartschat'


class HeadFinder:
    """Compute heads of mentions.

    This class provides functions to compute heads of mentions via modified
    version of the rules that can be found in Michael Collins' PhD thesis.
    The following changes were introduced:

        - handle NML as NP,
        - for coordinated phrases, take the coordination token as head,

    Furthermore, this class provides a function for adjusting heads for proper
    names to multi-token phrases via heuristics (see adjust_head_for_nam).
    """
    def __init__(self):
        self.__nonterminals = ["NP", "NML", "VP", "ADJP", "QP", "WHADVP", "S",
                             "ADVP", "WHNP", "SBAR", "SBARQ", "PP", "INTJ",
                             "SQ", "UCP", "X", "FRAG"]


        self.__nonterminal_rules = {
            "VP": (["TO", "VBD", "VBN", "MD", "VBZ", "VB", "VBG", "VBP", "VP",
                     "ADJP", "NN", "NNS", "NP"], False),
            "ADJP": (["NNS", "QP", "NN", "\$", "ADVP", "JJ", "VBN", "VBG", "ADJP",
              "JJR", "NP", "JJS", "DT", "FW", "RBR", "RBS", "SBAR", "RB"],
                     False),
            "QP": (["\$", "NNS", "NN", "IN", "JJ", "RB", "DT", "CD", "NCD",
              "QP", "JJR", "JJS"], False),
            "WHADVP": (["CC", "WRB"], True),
            "S": (["TO", "IN", "VP", "S", "SBAR", "ADJP", "UCP", "NP"], False),
            "SBAR": (["WHNP", "WHPP", "WHADVP", "WHADJP", "IN", "DT", "S", "SQ",
              "SINV", "SBAR", "FRAG"], False),
            "SBARQ": (["SQ", "S", "SINV", "SBARQ", "FRAG"], False),
            "SQ": (["VBZ", "VBD", "VBP", "VB", "MD", "VP", "SQ"], False),
            "ADVP": (["RB", "RBR", "RBS", "FW", "ADVP", "TO", "CD", "JJR", "JJ",
              "IN", "NP", "JJS", "NN"], True),
            "WHNP": (["WDT", "WP", "WP$", "WHADJP", "WHPP", "WHNP"], True),
            "PP": (["IN", "TO", "VBG", "VBN", "RP", "FW"], True),
            "X": (["S", "VP", "ADJP", "JJP", "NP", "SBAR", "PP", "X"], True),
            "FRAG": (["*"], True),
            "INTJ": (["*"], False),
            "UCP": (["*"], True),
        }

    def get_head(self, tree):
        """
        Compute the head of a mention, which is represented by its parse tree.

        Args:
            tree (nltk.ParentedTree): The parse tree of a mention.

        Returns:
            nltk.ParentedTree: The subtree of the input tree which corresponds
            to the head of the mention.
        """
        head = None

        label = tree.label()

        if len(tree) == 1:
            if tree.height() == 3:
                head = tree[0]
            elif tree.height() == 2:
                head = tree
        elif label in ["NP", "NML"]:
            head = self.__get_head_for_np(tree)
        elif label in self.__nonterminals:
            head = self.__get_head_for_nonterminal(tree)

        if head is None:
            head = self.get_head(tree[-1])

        return head

    def __get_head_for_np(self, tree):
        if self.__rule_cc(tree) is not None: