from abc import ABCMeta, abstractmethod, abstractproperty
from collections import defaultdict, Counter

import numpy as np
from numpy.testing import assert_approx_equal


def last_k(tokens, k):
    """Get the last k elements of a list as a tuple."""
    if not (0 <= k <= len(tokens)):
        raise ValueError('k must be between 0 and len(tokens) = {}, got: {}'.format(len(tokens), k))
    return tuple(tokens[len(tokens) - k:])


def replace_parens(tokens):
    """Replace all instances of -LRB- and -RRB- with actual parentheses."""
    parens_map = {'-LRB-': '(', '-RRB-': ')'}
    return [parens_map.get(s, s) for s in tokens]  # return identity if not parens symbols


def normalize_counts(counts):
    """Return a normalized Counter object."""
    normed = Counter()
    total = sum(counts.values(), 0.0)
    assert total > 0  # cannot normalize empty Counter
    for key, ct in counts.iteritems():
        normed[key] = ct / total
    normed.old_total = total  # document what the total was before normalization
    return normed


class LM(object):
    """Language model interface."""
    __metaclass__ = ABCMeta

    START = '<START>'
    END = '<END>'

    @abstractmethod
    def next_distribution(self, history):
        """Return a distribution over the next token.

        Args:
            history (List): a list of tokens generated so far

        Returns (Counter): a distribution
        """
        raise NotImplementedError

    @abstractproperty
    def max_context_size(self):
        """Return max allowed history context.

        Returns (int): maximum size of history context to keep
        """
        raise NotImplementedError


class CountLM(LM):
    """Naive language model.

    Uses empirical counts from the largest context it has observed. No sophisticated backoff strategy.

    Examples:
        lm = CountLM(4)

        # 'train' the language model
        for line in lines:
            tokens = line.split()
            lm.record_counts(tokens, append_end=True)
    """

    def __init__(self, max_context_size):
        """Construct a language model.

        Args:
            max_context_size (int): maximum # tokens to use as context
        """
        self._max_context_size = max_context_size
        self.contexts = defaultdict(Counter)

    @property
    def max_context_size(self):
        return self._max_context_size

    def _get_contexts(self, tokens):
        """List of contexts, from smallest to largest. Includes empty context.

        Returns:
            List[Tuple[str]]
        """
        contexts = []
        max_context = min(self._max_context_size, len(tokens))  # cannot go beyond max tokens
        for k in range(max_context + 1):
            contexts.append(last_k(tokens, k))
        return contexts

    def record_counts(self, tokens, append_end):
        """Record counts using `tokens` as a corpus.

        Args:
            tokens (List[str]): list of strings
        """
        history = [LM.START]
        if append_end:
            tokens = tokens + [LM.END]
        for tok in tokens:
            for context in self._get_contexts(history):
                self.contexts[context][tok] += 1
            history.append(tok)  # update history

    def _largest_context(self, history, contexts):
        """Find the largest context which matches history.

        Args:
            history (List[str]): a sequence of tokens
            contexts (Set[Tuple[str]]): a set of contexts, must include the empty context

        Returns:
            Tuple[str]: an item from contexts, which may be the empty context
        """
        assert tuple() in contexts  # empty context must be present

        for context in reversed(self._get_contexts(history)):
            if context in contexts:
                return context

    def _largest_known_context(self, history):
        """Find the largest recorded context which matches history."""
        return self._largest_context(history, self.contexts)

    def next_distribution(self, history):
        """Given a history, return a distribution (Counter) over the next token."""
        context = self._largest_known_context(history)
        counts = self.contexts[context]
        normed = normalize_counts(counts)
        normed.context = context
        return normed

    def sequence_probability(self, tokens):
        """Return the probability of each token in an article, based on the language model.

            Args:
                tokens (List): a list of tokens in the article

        Returns:
             List[Tuple[str, float]]: an ordered list of token-probability pairs"""
        history = [LM.START]
        probabilities = []

        for word in tokens:
            distr = self.next_distribution(history)
            if word in distr:
                probabilities.append((word, distr[word]))
            else:
                probabilities.append((word, 0.0))
            history.append(word)
        return probabilities


class KNNLM(LM):
    def __init__(self, article_embeddings, max_context_size, k_nearest):
        """Construct k-nearest-neighbor language model.

        Args:
            article_embeddings (ArticleEmbeddings): embeddings of each article
            max_context_size (int): max history to consider for CountLM
            k_nearest (int): # neighbors to consider
        """
        self.article_embeddings = article_embeddings
        self.k = k_nearest
        self.lm = CountLM(max_context_size)

    def record_nearest_counts(self, vec):
        name_score_pairs = self.article_embeddings.k_nearest_approx(vec, self.k)
        articles = [self.article_embeddings.name_to_article(name) for name, score in name_score_pairs]
        for art in articles:
            self.lm.record_counts(art.tokens, append_end=True)

    def next_distribution(self, history):
        return self.lm.next_distribution(history)

    @property
    def max_context_size(self):
        return self.lm.max_context_size

    def sequence_probability(self, tokens):
        return self.lm.sequence_probability(tokens)


class Generator(object):
    """Interface for language generator."""
    __metaclass__ = ABCMeta

    @abstractmethod
    def init_history(self):
        """Return a sequence of tokens to initialize the history."""
        pass

    @abstractmethod
    def get_next(self, history):
        """Get next token, given history."""
        pass

    @abstractmethod
    def stop_or_not(self, history):
        """Given what has been generated, decide whether to stop."""
        pass

    @abstractproperty
    def max_context_size(self):
        """Return max allowed history context.

        Returns (int): maximum size of history context to keep
        """
        raise NotImplementedError

    def truncate_history(self, history):
        """Truncate history when it grows much longer than max context size."""
        if len(history) > 2 * self.max_context_size:
            return list(last_k(history, self.max_context_size))
        return history

    def generate(self, history=None):
        """Generate a sequence of tokens."""
        if not history:
            history = self.init_history()
        return self.generate_custom(history, self.get_next, self.stop_or_not)

    def generate_custom(self, history, next_fxn, stop_fxn):
        """Generate a sequence using a custom next-token function and a custom stopping function.

        Args:
            history (List[T]): initial history
            next_fxn (Callable[[List[T]], T]): given a history, produce the next token
            stop_fxn (Callable[[List[T]], bool]): given a history, decide whether to stop
        """
        generated = []
        history = list(history)  # make a copy
        while True:
            next = next_fxn(history)
            history.append(next)
            history = self.truncate_history(history)
            if stop_fxn(history):
                break
            generated.append(next)
        return generated


class LMSampler(Generator):
    """Generation by sampling from a language model."""

    def __init__(self, lm):
        """Construct a LM sampler.

        Args:
            lm (LM): a language model
        """
        self.lm = lm

    @property
    def max_context_size(self):
        return self.lm.max_context_size

    def _sample_from_distribution(self, distr):
        """Sample from a categorical distribution.

        Args:
            distr (Counter): values must sum to 1
        Returns:
            one of the keys of distr
        """
        keys, probs = zip(*distr.items())
        assert_approx_equal(sum(probs), 1.)
        return np.random.choice(keys, p=probs)

    def init_history(self):
        return [self.lm.START]

    def get_next(self, history):
        return self._sample_from_distribution(self.lm.next_distribution(history))

    def stop_or_not(self, history):
        return history[-1] == LM.END

    @staticmethod
    def format_generation(tokens):
        return ' '.join(replace_parens(tokens))


class DistributionStats(object):
    def __init__(self, distr):
        self.total = distr.old_total
        self.context = distr.context

        probs = distr.values()
        assert_approx_equal(sum(probs), 1.)
        self.entropy = -1. * sum([p * np.log(p) for p in probs])

    def __repr__(self):
        return '{}:{}:{}'.format(len(self.context), self.total, self.entropy)


class LMSamplerWithStats(LMSampler):
    def init_history(self):
        return [(LM.START, 0)]

    def get_next(self, history):
        token_history, _ = zip(*history)
        distr = self.lm.next_distribution(token_history)
        next_token = self._sample_from_distribution(distr)
        return next_token, DistributionStats(distr)

    def stop_or_not(self, history):
        word = lambda pair: pair[0]
        return word(history[-1]) == LM.END

    @staticmethod
    def format_generation(token_stat_pairs):
        tokens, stats = zip(*list(token_stat_pairs))
        tokens = replace_parens(tokens)
        tokens = ['{:20}[{}]'.format(tok, stat) for tok, stat in zip(tokens, stats)]

        return '\n'.join(tokens)
