"""Statistical Language Processing tools.  (Chapter 22)
We define Unigram and Ngram text models, use them to generate random text,
and show the Viterbi algorithm for segmentatioon of letters into words.
Then we show a very simple Information Retrieval system, and an example
working on a tiny sample of Unix manual pages."""

from utils import argmin
from learning import CountingProbDist
import search

from math import log, exp
from collections import defaultdict
import heapq
import re
import os


class UnigramTextModel(CountingProbDist):

    """This is a discrete probability distribution over words, so you
    can add, sample, or get P[word], just like with CountingProbDist.  You can
    also generate a random text n words long with P.samples(n)"""

    def samples(self, n):
        "Return a string of n words, random according to the model."
        return ' '.join(self.sample() for i in range(n))



class NgramTextModel(CountingProbDist):

    """This is a discrete probability distribution over n-tuples of words.
    You can add, sample or get P[(word1, ..., wordn)]. The method P.samples(n)
    builds up an n-word sequence; P.add and P.add_sequence add data."""

    def __init__(self, n, observation_sequence=[], default=0):
        # In addition to the dictionary of n-tuples, cond_prob is a
        # mapping from (w1, ..., wn-1) to P(wn | w1, ... wn-1)
        CountingProbDist.__init__(self, default=default)
        self.n = n
        self.cond_prob = defaultdict()
        self.add_sequence(observation_sequence)

    # __getitem__, top, sample inherited from CountingProbDist
    # Note they deal with tuples, not strings, as inputs

    def add(self, ngram):
        """Count 1 for P[(w1, ..., wn)] and for P(wn | (w1, ..., wn-1)"""
        CountingProbDist.add(self, ngram)
        if ngram[:-1] not in self.cond_prob:
            self.cond_prob[ngram[:-1]] = CountingProbDist()
        self.cond_prob[ngram[:-1]].add(ngram[-1])

    def add_empty(self, words, n):
        return [''] * (n - 1) + words

    def add_sequence(self, words):
        """Add each of the tuple words[i:i+n], using a sliding window.
        Prefix some copies of the empty word, '', to make the start work."""
        n = self.n
        words = self.add_empty(words, n)

        for i in range(len(words) - n):
            self.add(tuple(words[i:i + n]))

    def samples(self, nwords):
        """Build up a random sample of text nwords words long, using
        the conditional probability given the n-1 preceding words."""
        n = self.n
        nminus1gram = ('',) * (n-1)
        output = []
        for i in range(nwords):
            if nminus1gram not in self.cond_prob:
                nminus1gram = ('',) * (n-1)  # Cannot continue, so restart.
            wn = self.cond_prob[nminus1gram].sample()
            output.append(wn)
            nminus1gram = nminus1gram[1:] + (wn,)
        return ' '.join(output)


class NgramCharModel(NgramTextModel):
    def add_empty(self, words, n):
        return  ' ' * (n - 1) + words

    def add_sequence(self, words):
        for word in words:
            super().add_sequence(word)

# ______________________________________________________________________________


def viterbi_segment(text, P):
    """Find the best segmentation of the string of characters, given the
    UnigramTextModel P."""
    # best[i] = best probability for text[0:i]
    # words[i] = best word ending at position i
    n = len(text)
    words = [''] + list(text)
    best = [1.0] + [0.0] * n
    # Fill in the vectors best, words via dynamic programming
    for i in range(n+1):
        for j in range(0, i):
            w = text[j:i]
            if P[w] * best[i - len(w)] >= best[i]:
                best[i] = P[w] * best[i - len(w)]
                words[i] = w
    # Now recover the sequence of best words
    sequence = []
    i = len(words) - 1
    while i > 0:
        sequence[0:0] = [words[i]]
        i = i - len(words[i])
    # Return sequence of best words and overall probability
    return sequence, best[-1]


# ______________________________________________________________________________


# TODO(tmrts): Expose raw index
class IRSystem:

    """A very simple Information Retrieval System, as discussed in Sect. 23.2.
    The constructor s = IRSystem('the a') builds an empty system with two
    stopwords. Next, index several documents with s.index_document(text, url).
    Then ask queries with s.query('query words', n) to retrieve the top n
    matching documents.  Queries are literal words from the document,
    except that stopwords are ignored, and there is one special syntax:
    The query "learn: man cat", for example, runs "man cat" and indexes it."""

    def __init__(self, stopwords='the a of'):
        """Create an IR System. Optionally specify stopwords."""
        # index is a map of {word: {docid: count}}, where docid is an int,
        # indicating the index into the documents list.
        self.index = defaultdict(lambda: defaultdict(int))
        self.stopwords = set(words(stopwords))
        self.documents = []

    def index_collection(self, filenames):
        "Index a whole collection of files."
        prefix = os.path.dirname(__file__)
        for filename in filenames:
            self.index_document(open(filename).read(),
                                os.path.relpath(filename, prefix))

    def index_document(self, text, url):
        "Index the text of a document."
        # For now, use first line for title
        title = text[:text.index('\n')].strip()
        docwords = words(text)
        docid = len(self.documents)
        self.documents.append(Document(title, url, len(docwords)))
        for word in docwords:
            if word not in self.stopwords:
                self.index[word][docid] += 1

    def query(self, query_text, n=10):
        """Return a list of n (score, docid) pairs for the best matches.
        Also handle the special syntax for 'learn: command'."""
        if query_text.startswith("learn:"):
            doctext = os.popen(query_text[len("learn:"):], 'r').read()
            self.index_document(doctext, query_text)
            return []
        qwords = [w for w in words(query_text) if w not in self.stopwords]