#===============================================================================
# Copyright 2008 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.
#===============================================================================

"""
The highlight module contains classes and functions for displaying short
excerpts from hit documents in the search results you present to the user, with
query terms highlighted.
"""

from __future__ import division
from heapq import nlargest

# Fragment object

class Fragment(object):
    """Represents a fragment (extract) from a hit document. This object is mainly
    used to keep track of the start and end points of the fragment; it does not
    contain the text of the fragment or do much else.
    """
    
    def __init__(self, tokens, charsbefore = 0, charsafter = 0, textlen = 999999):
        """
        :param tokens: list of the Token objects representing the matched terms. 
        :param charsbefore: approx. how many characters before the start of the first
            matched term to include in the fragment.
        :param charsafter: approx. how many characters after the end of the last
            matched term to include in the fragment.
        :param textlen: length in characters of the document text.
        """
        
        #: index of the first character of the fragment in the original document
        self.startchar = max(0, tokens[0].startchar - charsbefore)
        #: index after the last character of the fragment in the original document
        self.endchar = min(textlen, tokens[-1].endchar + charsafter)
        self.matches = [t for t in tokens if t.matched]
        self.matched_terms = frozenset(t.text for t in self.matches)
    
    def __len__(self):
        return self.endchar - self.startchar
    
    def overlaps(self, fragment):
        sc = self.startchar
        ec = self.endchar
        fsc = fragment.startchar
        fec = fragment.endchar
        return (fsc > sc and fsc < ec) or (fec > sc and fec < ec)
    
    def overlapped_length(self, fragment):
        sc = self.startchar
        ec = self.endchar
        fsc = fragment.startchar
        fec = fragment.endchar
        return max(ec, fec) - min(sc, fsc)
    
    def has_matches(self):
        return any(t.matched for t in self.tokens)
    

# Filters

def copyandmatchfilter(termset, tokens):
    for t in tokens:
        t = t.copy()
        t.matched = t.text in termset
        yield t

# Fragmenters

def NullFragmenter(text, termset, tokens):
    """Doesn't fragment the token stream. This object just
    returns the entire stream as one "fragment". This is useful if
    you want to highlight the entire text.
    """
    return [Fragment(list(tokens))]


class SimpleFragmenter(object):
    """Simply splits the text into roughly equal sized chunks.
    """
    
    def __init__(self, size = 70):
        """
        :param size: size (in characters) to chunk to. The chunking is based on
            tokens, so the fragments will usually be smaller.
        """
        self.size = 70
        
    def __call__(self, text, tokens):
        size = self.size
        first = None
        frag = []
        
        for t in tokens:
            if first is None:
                first = t.startchar
            
            if t.endchar - first > size:
                first = None
                if frag:
                    yield Fragment(frag)
                frag = []
            
            frag.append(t)
            
        if frag:
            yield Fragment(frag)


class SentenceFragmenter(object):
    """Breaks the text up on sentence end punctuation characters (".", "!", or "?").
    This object works by looking in the original text for a sentence end as the next
    character after each token's 'endchar'.
    """
    
    def __init__(self, maxchars = 200, sentencechars = ".!?"):
        """
        :param maxchars: The maximum number of characters allowed in a fragment.
        """
        
        self.maxchars = maxchars
        self.sentencechars = frozenset(sentencechars)
    
    def __call__(self, text, tokens):
        maxchars = self.maxchars
        sentencechars = self.sentencechars
        textlen = len(text)
        first = None
        frag = []
        
        for t in tokens:
            if first is None:
                first = t.startchar
            endchar = t.endchar
            
            if endchar - first > maxchars:
                first = None
                if frag:
                    yield Fragment(frag)
                frag = []
            
            frag.append(t)
            if frag and endchar < textlen and text[endchar] in sentencechars:
                # Don't break for two periods in a row (e.g. ignore "...")
                if endchar+1 < textlen and text[endchar + 1] in sentencechars:
                    continue
                
                yield Fragment(frag, charsafter = 1)
                frag = []
                first = None
        
        if frag:
            yield Fragment(frag)


class ContextFragmenter(object):
    """Looks for matched terms and aggregates them with their
    surrounding context.
    
    This fragmenter only yields fragments that contain matched terms.    
    """
    
    def __init__(self, termset, maxchars = 200, charsbefore = 20, charsafter = 20):
        """
        :param termset: A collection (probably a set or frozenset) containing the
            terms you want to match to token.text attributes.
        :param maxchars: The maximum number of characters allowed in a fragment.
        :param charsbefore: The number of extra characters of context to add before
            the first matched term.
        :param charsafter: The number of extra characters of context to add after
            the last matched term.
        """
        
        self.maxchars = maxchars
        self.charsbefore = charsbefore
        self.charsafter = charsafter
        
    def __call__(self, text, tokens):
        maxchars = self.maxchars
        charsbefore = self.charsbefore
        charsafter = self.charsafter
        
        current = []
        currentlen = 0
        countdown = -1
        for t in tokens:
            if t.matched:
                countdown = charsafter
            
            current.append(t)
            
            length = t.endchar - t.startchar
            currentlen += length
            
            if countdown >= 0:
                countdown -= length
                
                if countdown < 0 or currentlen >= maxchars:
                    yield Fragment(current)
                    current = []
                    currentlen = 0
            
            else:
                while current and currentlen > charsbefore:
                    t = current.pop(0)
                    currentlen -= t.endchar - t.startchar

        if countdown >= 0:
            yield Fragment(current)


#class VectorFragmenter(object):
#    def __init__(self, termmap, maxchars = 200, charsbefore = 20, charsafter = 20):
#        """
#        :param termmap: A dictionary mapping the terms you're looking for to
#            lists of either (posn, startchar, endchar) or
#            (posn, startchar, endchar, boost) tuples.
#        :param maxchars: The maximum number of characters allowed in a fragment.
#        :param charsbefore: The number of extra characters of context to add before
#            the first matched term.
#        :param charsafter: The number of extra characters of context to add after
#            the last matched term.
#        """
#        
#        self.termmap = termmap
#        self.maxchars = maxchars
#        self.charsbefore = charsbefore
#        self.charsafter = charsafter
#    
#    def __call__(self, text, tokens):
#        maxchars = self.maxchars
#        charsbefore = self.charsbefore
#        charsafter = self.charsafter
#        textlen = len(text)
#        
#        vfrags = []
#        for term, data in self.termmap.iteritems():
#            if len(data) == 3:
#                t = Token(startchar = data[1], endchar = data[2])
#            elif len(data) == 4:
#                t = Token(startchar = data[1], endchar = data[2], boost = data[3])
#            else:
#                raise ValueError(repr(data))
#            
#            newfrag = VFragment([t], charsbefore, charsafter, textlen)
#            added = False
#            
#            for vf in vfrags:
#                if vf.overlaps(newfrag) and vf.overlapped_length(newfrag) < maxchars:
#                    vf.merge(newfrag)
#                    added = True
#                    break


# Fragment scorers

def BasicFragmentScorer(f):
    # Add up the boosts for the matched terms in this passage
    score = sum(t.boost for t in f.matches)
    
    # Favor diversity: multiply score by the number of separate
    # terms matched
    score *= len(f.matched_terms) * 100
    
    return score


# Fragment sorters

def SCORE(fragment):
    "Sorts higher scored passages first."
    return None
def FIRST(fragment):
    "Sorts passages from earlier in the document first."
    return fragment.startchar
def LONGER(fragment):
    "Sorts longer passages first."
    return 0 - len(fragment)
def SHORTER(fragment):
    "Sort shorter passages first."
    return len(fragment)


# Formatters

class UppercaseFormatter(object):
    def __init__(self, between = "..."):
        self.between = between
        
    def _format_fragment(self, text, fragment):
        output = []
        index = fragment.startchar
        
        for t in fragment.matches:
            if t.startchar > index:
                output.append(text[index:t.startchar])
            
            ttxt = text[t.startchar:t.endchar]
            if t.matched: ttxt = ttxt.upper()
            output.append(ttxt)
            index = t.endchar
        
        output.append(text[index:fragment.endchar])
        return "".join(output)

    def __call__(self, text, fragments):
        return self.between.join((self._format_fragment(text, fragment)
                                  for fragment in fragments))


class GenshiFormatter(object):
    def __init__(self, qname, between = "..."):
        self.qname = qname
        self.between = between
        
        from genshi.core import START, END, TEXT, Attrs, Stream #@UnresolvedImport
        self.START, self.END, self.TEXT, self.Attrs, self.Stream = (START, END, TEXT, Attrs, Stream)

    def _add_text(self, text, output):
        if output and output[-1][0] == self.TEXT:
            output[-1] = (self.TEXT, output[-1][1] + text, output[-1][2])
        else:
            output.append((self.TEXT, text, (None, -1, -1)))

    def _format_fragment(self, text, fragment):
        START, TEXT, END, Attrs = self.START, self.TEXT, self.END, self.Attrs
        qname = self.qname
        output = []
        
        index = fragment.startchar
        lastmatched = False
        for t in fragment.matches:
            if t.startchar > index:
                if lastmatched:
                    output.append((END, qname, (None, -1, -1)))
                    lastmatched = False
                self._add_text(text[index:t.startchar], output)
            
            ttxt = text[t.startchar:t.endchar]
            if not lastmatched:
                output.append((START, (qname, Attrs()), (None, -1, -1)))
                lastmatched = True
            output.append((TEXT, ttxt, (None, -1, -1)))
                                    
            index = t.endchar
        
        if lastmatched:
            output.append((END, qname, (None, -1, -1)))
        
        return output

    def __call__(self, text, fragments):
        output = []
        first = True
        for fragment in fragments:
            if not first:
                self._add_text(self.between, output)
            first = False
            output += self._format_fragment(text, fragment)
        
        return self.Stream(output)


# Highlighting

def top_fragments(text, terms, analyzer, fragmenter, top = 3,
                  scorer = BasicFragmentScorer, minscore = 1):
    termset = frozenset(terms)
    tokens = copyandmatchfilter(termset,
                                analyzer(text, chars = True, keeporiginal = True))
    
    scored_frags = nlargest(top, ((scorer(f), f) for f in fragmenter(text, tokens)))
    return [sf for score, sf in scored_frags if score > minscore]


def highlight(text, terms, analyzer, fragmenter, formatter, top=3,
              scorer = BasicFragmentScorer, minscore = 1,
              order = FIRST):
    
    fragments = top_fragments(text, terms, analyzer, fragmenter,
                              top = top, minscore = minscore)
    fragments.sort(key = order)
    return formatter(text, fragments)
    

if __name__ == '__main__':
    import re, time
    from whoosh import analysis
    #from genshi import QName
    
    sa = analysis.StemmingAnalyzer()
    txt = open("/Volumes/Storage/Development/help/documents/nodes/sop/copy.txt").read().decode("utf8")
    txt = re.sub("[\t\r\n ]+", " ", txt)
    t = time.time()
    fs = highlight(txt, ["templat", "geometri"], sa, SentenceFragmenter(), UppercaseFormatter())
    print time.time() - t
    print fs




