import logging

logger = logging.getLogger(None if __name__ == '__main__' else __name__)


class ConstituencyTree:
    """
    Group together information from a constituency tree and expose convenient
    lookups.
    """

    def __init__(self, head2deps):
        """
        Initialise `ConstituencyTree`
        """
        self.head2deps = head2deps

        # Create the reverse too
        self.dep2heads = {}
        for headID, deps in head2deps.items():
            for toID, relation in deps:
                self.dep2heads.setdefault(toID, set()).add((headID, relation))

        logger.debug(self.head2deps, self.dep2heads)
        self._head2constituent = {}

    def __repr__(self):
        return self.__class__.__name__ + '({!r})'.format(self.head2deps)

    @classmethod
    def from_naf(cls, nafobj, term_filter=lambda naf, t: True):
        """
        Initialise this ConstituencyTree from a NAF object

        Only keep terms where `term_filter(term)` evaluates True.

        :param nafobj:          NAF object from input
        :param term_filter:     filter for terms
        """
        return cls(cls.create_headdep_dict(
            nafobj,
            term_filter
        ))

    def get_direct_dependents(self, ID):
        """
        Get the term IDs of the terms directly dependent on `ID`.

        :param headID:  term ID to get direct dependents of
        :return:        a set of IDs that are direct dependents of `ID`
        """
        with_rel = self.get_direct_dependents_with_relation(ID)
        return {ID for ID, _ in with_rel} if with_rel is not None else None

    def get_direct_dependents_with_relation(self, ID):
        """
        Get the (ID, relation) tuples of the terms directly dependent on `ID`.

        :param headID:  term ID to get direct dependents of
        :return:        a set (ID, relation) tuples that are direct dependents
                        of `ID`
        """
        return self.head2deps.get(ID, None)

    def get_constituent(self, ID):
        """
        Get the term IDs of the terms dependent on `ID`.

        Always contains `ID`, even if it is an unknown `ID`.

        :param headID:  term ID to get constituent of
        :return:        a set of IDs of terms that are dependents of `ID`
        """
        try:
            return self._head2constituent[ID]
        except KeyError:
            return self._get_constituent(ID, set())

    def _get_constituent(self, ID, parents):
        """
        Calculate the constituent if it is not in the cache

        :param head:      head to find constituent of
        :param parents:   term IDs that are already being calculated, thus
                          should not be recursed into
        :return:          set of term IDs that are in the constituent of `head`
        """
        if ID in parents:
            # Prevent loop
            return set()
        elif ID in self._head2constituent:
            # Use cache
            return self._head2constituent[ID]
        # Base case: direct dependents
        deps = {term_id for term_id, _ in self.head2deps.get(ID, [])}
        # Recursive step: call _get_constituent for every direct dependent
        deps.union(*(
            self._get_constituent(term_id, parents | {term_id})
            for term_id in deps
        ))
        # Make sure ID itself is in it
        deps.add(ID)
        # Cache
        self._head2constituent[ID] = deps
        return deps

    def get_direct_parents(self, ID):
        """
        Get the term IDs of the terms of which `ID` is a direct dependent.

        :param headID:  term ID to get direct dependents of
        :return:        a set of IDs that are direct dependents of `ID`
        """
        raise NotImplementedError()

    @staticmethod
    def create_headdep_dict(nafobj, term_filter):
        """
        Create dictionaries of dependent to heads and head to direct dependent

        Only keep terms where `term_filter(term)` evaluates True.

        :param nafobj:          NAF object from input
        :param term_filter:     filter for terms
        :return:        head2deps, dep2heads
        """

        allhead2deps = {}
        dep2headIDs = {}
        for dep in nafobj.get_dependencies():
            headID = dep.get_from()
            toID = dep.get_to()
            allhead2deps.setdefault(headID, []).append(
                (toID, dep.get_function())
            )
            dep2headIDs.setdefault(toID, []).append(headID)

        head2deps = {}
        for headID, deps in allhead2deps.items():
            # I don't have to do something with the deps that are filtered out,
            # because if they are leaves they can just be deleted and if they
            # aren't leaves they will also appear as headID and handled there.
            deps = {
                (toID, relation)
                for toID, relation in deps
                if term_filter(nafobj, toID)
            }
            if term_filter(nafobj, headID):
                head2deps.setdefault(headID, set()).update(deps)
            else:
                # Delete the head by adding its dependents to the heads of the
                # head.
                for super_headID in dep2headIDs[headID]:
                    if term_filter(nafobj, super_headID):
                        head2deps.setdefault(super_headID, set()).update(deps)

        return head2deps

@@ -1476,267 +1476,203 @@
     
-with_rel = self.get_direct_dependents_with_relation(ID)%0A        return %7BID for ID, _ in with_rel%7D if with_rel is not None else None%0A%0A    def get_direct_dependents_with_relation(self, ID):%0A        %22%22%22%0A        Get the (ID, relation) tuples of the terms
+if ID in self.head2deps:%0A            return %7BID for ID, _ in self.head2deps%5BID%5D%7D%0A%0A    def get_direct_parents(self, ID):%0A        %22%22%22%0A        Get the term IDs of the terms of which %60ID%60 is a
  direct
-ly
  dep
@@ -1673,32 +1673,24 @@
 ct dependent
- on %60ID%60
 .%0A%0A        :
@@ -1718,38 +1718,35 @@
 D to get direct 
-depend
+par
 ents of%0A        
@@ -1771,28 +1771,13 @@
 set 
-(ID, relation) tuple
+of ID
 s th
@@ -1794,29 +1794,32 @@
 ect 
-dependents%0A  
+parents of %60ID%60%0A
         
     
@@ -1810,32 +1810,36 @@
 ID%60%0A        
+%22%22%22%0A
         
       of %60ID
@@ -1830,77 +1830,88 @@
     
-      o
+i
 f 
-%60
 ID
-%60%0A        %22%22%22%0A        return self.head2deps.get(ID, None)
+ in self.dep2heads:%0A            return %7BID for ID, _ in self.dep2heads%5BID%5D%7D
 %0A%0A  
@@ -3392,315 +3392,8 @@
 ps%0A%0A
-    def get_direct_parents(self, ID):%0A        %22%22%22%0A        Get the term IDs of the terms of which %60ID%60 is a direct dependent.%0A%0A        :param headID:  term ID to get direct dependents of%0A        :return:        a set of IDs that are direct dependents of %60ID%60%0A        %22%22%22%0A        raise NotImplementedError()%0A%0A
     
@@ -4322,23 +4322,19 @@
 can 
-just be deleted
+be left out
  and

Write `get_direct_parents`