#===============================================================================
# Copyright 2009 Matt Chaput
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#    http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#===============================================================================

"""This module contains low-level functions and a high-level class for parsing
the prolog file "wn_s.pl" from the WordNet prolog download
into an object suitable for looking up synonyms and performing query expansion.

http://wordnetcode.princeton.edu/3.0/WNprolog-3.0.tar.gz
"""

from collections import defaultdict

from whoosh.fields import Schema, ID, STORED
from whoosh.index import Index


def parse_file(f):
    """Parses the WordNet wn_s.pl prolog file and returns two dictionaries:
    word2nums and num2words.
    """
    
    word2nums = defaultdict(list)
    num2words = defaultdict(list)
    
    for line in f:
        if not line.startswith("s("):
            continue
        
        line = line[2:]
        num = int(line[:line.find(",")])
        qt = line.find("'")
        line = line[qt+1:]
        qt = line.find("'")
        word = line[:qt].lower()

        if not word.isalpha():
            continue
        
        word2nums[word].append(num)
        num2words[num].append(word)
    
    return word2nums, num2words


def make_index(storage, indexname, word2nums, num2words):
    """Creates a Whoosh index in the given storage object containing
    synonyms taken from word2nums and num2words. Returns the Index
    object.
    """
    
    schema = Schema(word=ID, syns=STORED)
    ix = storage.create_index(schema, indexname=indexname)
    w = ix.writer()
    for word in word2nums.iterkeys():
        syns = synonyms(word2nums, num2words, word)
        w.add_document(word=unicode(word), syns=syns)
    w.commit()
    return ix


def synonyms(word2nums, num2words, word):
    """Uses the word2nums and num2words dicts to look up synonyms
    for the given word. Returns a list of synonym strings.
    """
    
    keys = word2nums[word]
    syns = set()
    for key in keys:
        syns = syns.union(num2words[key])
    
    if word in syns:
        syns.remove(word)
    return sorted(syns)


class Thesaurus(object):
    """Represents the WordNet synonym database, either loaded into memory
    from the wn_s.pl Prolog file, or stored on disk in a Whoosh index.
    
    This class allows you to parse the prolog file "wn_s.pl" from the WordNet prolog
    download into an object suitable for looking up synonyms and performing query
    expansion.

    http://wordnetcode.princeton.edu/3.0/WNprolog-3.0.tar.gz
    
    To load a Thesaurus object from the wn_s.pl file...
    
    >>> t = Thesaurus.from_filename("wn_s.pl")
    
    To save the in-memory Thesaurus to a Whoosh index...
    
    >>> from whoosh.filedb.filestore import FileStorage
    >>> fs = FileStorage("index")
    >>> t.to_storage(fs)
    
    To load a Thesaurus object from a Whoosh index...
    
    >>> t = Thesaurus.from_storage(fs)
    
    The Thesaurus object is thus usable in two ways:
    
    * Parse the wn_s.pl file into memory (Thesaurus.from_*) and then look up
      synonyms in memory. This has a startup cost for parsing the file, and uses
      quite a bit of memory to store two large dictionaries, however synonym
      look-ups are very fast.
      
    * Parse the wn_s.pl file into memory (Thesaurus.from_filename) then save it to
      an index (to_storage). From then on, open the thesaurus from the saved
      index (Thesaurus.from_storage). This has a large cost for storing the index,
      but after that it is faster to open the Thesaurus (than re-parsing the file)
      but slightly slower to look up synonyms.
    
    Here are timings for various tasks on my (fast) Windows machine, which might
    give an idea of relative costs for in-memory vs. on-disk.
    
    ================================================ ================
    Task                                             Approx. time (s)
    ================================================ ================
    Parsing the wn_s.pl file                         1.045
    Saving to an on-disk index                       13.084
    Loading from an on-disk index                    0.082
    Look up synonyms for "light" (in memory)         0.0011
    Look up synonyms for "light" (loaded from disk)  0.0028
    ================================================ ================
    
    Basically, if you can afford spending the memory necessary to parse the
    Thesaurus and then cache it, it's faster. Otherwise, use an on-disk index.
    """
    
    def __init__(self):
        self.w2n = None
        self.n2w = None
        self.searcher = None
    
    @classmethod
    def from_file(cls, fileobj):
        """Creates a Thesaurus object from the given file-like object, which should
        contain the WordNet wn_s.pl file.
        
        >>> f = open("wn_s.pl")
        >>> t = Thesaurus.from_file(f)
        >>> t.synonyms("hail")
        ['acclaim', 'come', 'herald']
        """
        
        thes = cls()
        thes.w2n, thes.n2w = parse_file(fileobj)
        return thes
    
    @classmethod
    def from_filename(cls, filename):
        """Creates a Thesaurus object from the given filename, which should
        contain the WordNet wn_s.pl file.
        
        >>> t = Thesaurus.from_filename("wn_s.pl")
        >>> t.synonyms("hail")
        ['acclaim', 'come', 'herald']
        """
        
        f = open(filename, "rb")
        try:
            return cls.from_file(f)
        finally:
            f.close()
    
    @classmethod
    def from_storage(cls, storage, indexname="THES"):
        """Creates a Thesaurus object from the given storage object,
        which should contain an index created by Thesaurus.to_storage().
        
        >>> from whoosh.filedb.filestore import FileStorage
        >>> fs = FileStorage("index")
        >>> t = Thesaurus.from_storage(fs)
        >>> t.synonyms("hail")
        ['acclaim', 'come', 'herald']
        
        :param storage: A :class:`whoosh.store.Storage` object from
            which to load the index.
        :param indexname: A name for the index. This allows you to
            store multiple indexes in the same storage object.
        """
        
        thes = cls()
        index = storage.open_index(indexname=indexname)
        thes.searcher = index.searcher()
        return thes
    
    def to_storage(self, storage, indexname="THES"):
        """Creates am index in the given storage object from the
        synonyms loaded from a WordNet file.
        
        >>> from whoosh.filedb.filestore import FileStorage
        >>> fs = FileStorage("index")
        >>> t = Thesaurus.from_filename("wn_s.pl")
        >>> t.to_storage(fs)
        
        :param storage: A :class:`whoosh.store.Storage` object in
            which to save the index.
        :param indexname: A name for the index. This allows you to
            store multiple indexes in the same storage object.
        """
        
        if not self.w2n or not self.n2w:
            raise Exception("No synonyms loaded")
        make_index(storage, indexname, self.w2n, self.n2w)

    def synonyms(self, word):
        """Returns a list of synonyms for the given word.
        
        >>> thesaurus.synonyms("hail")
        ['acclaim', 'come', 'herald']
        """
        
        word = word.lower()
        if self.searcher:
            return self.searcher.document(word=word)["syns"]
        else:
            return synonyms(self.w2n, self.n2w, word)


if __name__ == "__main__":
    from time import clock
    from whoosh.filedb.filestore import FileStorage
    st = FileStorage("c:/testindex")
    
#    t = clock()
#    th = Thesaurus.from_filename("c:/wordnet/wn_s.pl")
#    print clock() - t
#    
#    t = clock()
#    th.to_storage(st)
#    print clock() - t
#    
#    t = clock()
#    print th.synonyms("light")
#    print clock() - t
    
    t = clock()
    th = Thesaurus.from_storage(st)
    print clock() - t
    
    t = clock()
    print th.synonyms("hail")
    print clock() - t
    
    