# Natural Language Toolkit: Verbnet Corpus Reader
#
# Copyright (C) 2001-2013 NLTK Project
# Author: Edward Loper <edloper@gradient.cis.upenn.edu>
# URL: <http://www.nltk.org/>
# For license information, see LICENSE.TXT
from __future__ import unicode_literals

import re
import textwrap
from collections import defaultdict

from nltk import compat
from .xmldocs import XMLCorpusReader

class VerbnetCorpusReader(XMLCorpusReader):

    # No unicode encoding param, since the data files are all XML.
    def __init__(self, root, fileids, wrap_etree=False):
        XMLCorpusReader.__init__(self, root, fileids, wrap_etree)

        self._lemma_to_class = defaultdict(list)
        """A dictionary mapping from verb lemma strings to lists of
        verbnet class identifiers."""

        self._wordnet_to_class = defaultdict(list)
        """A dictionary mapping from wordnet identifier strings to
        lists of verbnet class identifiers."""

        self._class_to_fileid = {}
        """A dictionary mapping from class identifiers to
        corresponding file identifiers.  The keys of this dictionary
        provide a complete list of all classes and subclasses."""

        self._shortid_to_longid = {}

        # Initialize the dictionaries.  Use the quick (regexp-based)
        # method instead of the slow (xml-based) method, because it
        # runs 2-30 times faster.
        self._quick_index()

    _LONGID_RE = re.compile(r'([^\-\.]*)-([\d+.\-]+)$')
    """Regular expression that matches (and decomposes) longids"""

    _SHORTID_RE = re.compile(r'[\d+.\-]+$')
    """Regular expression that matches shortids"""

    _INDEX_RE = re.compile(r'<MEMBER name="\??([^"]+)" wn="([^"]*)"[^>]+>|'
                           r'<VNSUBCLASS ID="([^"]+)"/?>')
    """Regular expression used by ``_index()`` to quickly scan the corpus
       for basic information."""

    def lemmas(self, classid=None):
        """
        Return a list of all verb lemmas that appear in any class, or
        in the ``classid`` if specified.
        """
        if classid is None:
            return sorted(self._lemma_to_class.keys())
        else:
            # [xx] should this include subclass members?
            vnclass = self.vnclass(classid)
            return [member.get('name') for member in
                    vnclass.findall('MEMBERS/MEMBER')]

    def wordnetids(self, classid=None):
        """
        Return a list of all wordnet identifiers that appear in any
        class, or in ``classid`` if specified.
        """
        if classid is None:
            return sorted(self._wordnet_to_class.keys())
        else:
            # [xx] should this include subclass members?
            vnclass = self.vnclass(classid)
            return sum([member.get('wn','').split() for member in
                        vnclass.findall('MEMBERS/MEMBER')], [])

    def classids(self, lemma=None, wordnetid=None, fileid=None, classid=None):
        """
        Return a list of the verbnet class identifiers.  If a file
        identifier is specified, then return only the verbnet class
        identifiers for classes (and subclasses) defined by that file.
        If a lemma is specified, then return only verbnet class
        identifiers for classes that contain that lemma as a member.
        If a wordnetid is specified, then return only identifiers for
        classes that contain that wordnetid as a member.  If a classid
        is specified, then return only identifiers for subclasses of
        the specified verbnet class.
        """
        if len([x for x in [lemma, wordnetid, fileid, classid]
                if x is not None]) > 1:
            raise ValueError('Specify at most one of: fileid, wordnetid, '
                             'fileid, classid')
        if fileid is not None:
            return [c for (c,f) in self._class_to_fileid.items()
                    if f == fileid]
        elif lemma is not None:
            return self._lemma_to_class[lemma]
        elif wordnetid is not None:
            return self._wordnet_to_class[wordnetid]
        elif classid is not None:
            xmltree = self.vnclass(classid)
            return [subclass.get('ID') for subclass in
                    xmltree.findall('SUBCLASSES/VNSUBCLASS')]
        else:
            return sorted(self._class_to_fileid.keys())

    def vnclass(self, fileid_or_classid):
        """
        Return an ElementTree containing the xml for the specified
        verbnet class.

        :param fileid_or_classid: An identifier specifying which class
            should be returned.  Can be a file identifier (such as
            ``'put-9.1.xml'``), or a verbnet class identifier (such as
            ``'put-9.1'``) or a short verbnet class identifier (such as
            ``'9.1'``).
        """
        # File identifier: just return the xml.
        if fileid_or_classid in self._fileids:
            return self.xml(fileid_or_classid)

        # Class identifier: get the xml, and find the right elt.
        classid = self.longid(fileid_or_classid)
        if classid in self._class_to_fileid:
            fileid = self._class_to_fileid[self.longid(classid)]
            tree = self.xml(fileid)
            if classid == tree.get('ID'):
                return tree
            else:
                for subclass in tree.findall('.//VNSUBCLASS'):
                    if classid == subclass.get('ID'):
                        return subclass
                else:
                    assert False # we saw it during _index()!

        else:
            raise ValueError('Unknown identifier %s' % fileid_or_classid)

    def fileids(self, vnclass_ids=None):
        """
        Return a list of fileids that make up this corpus.  If
        ``vnclass_ids`` is specified, then return the fileids that make
        up the specified verbnet class(es).
        """
        if vnclass_ids is None: