# Ported from ref/word_cloud/wordcloud/tokenization.py (MIT License)
from __future__ import division
from itertools import tee
from operator import itemgetter
from collections import defaultdict
from math import log


def _l(k, n, x):
    """Dunning log-likelihood helper."""
    return log(max(x, 1e-10)) * k + log(max(1 - x, 1e-10)) * (n - k)


def _collocation_score(count_bigram, count1, count2, n_words):
    """Dunning likelihood ratio collocation score."""
    if n_words <= count1 or n_words <= count2:
        return 0
    N, c12, c1, c2 = n_words, count_bigram, count1, count2
    p = c2 / N
    p1 = c12 / c1
    p2 = (c2 - c12) / (N - c1)
    score = (
        _l(c12, c1, p) + _l(c2 - c12, N - c1, p)
        - _l(c12, c1, p1) - _l(c2 - c12, N - c1, p2)
    )
    return -2 * score


def _pairwise(iterable):
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)


def process_tokens(words, normalize_plurals=True):
    """Count words, normalizing case and optionally merging plurals.

    Returns
    -------
    counts : dict str -> int
    standard_forms : dict lowercase_str -> canonical_str
    """
    d = defaultdict(dict)
    for word in words:
        wl = word.lower()
        case_dict = d[wl]
        case_dict[word] = case_dict.get(word, 0) + 1

    if normalize_plurals:
        merged_plurals = {}
        for key in list(d.keys()):
            if key.endswith('s') and not key.endswith('ss'):
                singular = key[:-1]
                if singular in d:
                    for word, count in d[key].items():
                        sing_form = word[:-1]
                        d[singular][sing_form] = d[singular].get(sing_form, 0) + count
                    merged_plurals[key] = singular
                    del d[key]

    fused_cases = {}
    standard_cases = {}
    item1 = itemgetter(1)
    for word_lower, case_dict in d.items():
        first = max(case_dict.items(), key=item1)[0]
        fused_cases[first] = sum(case_dict.values())
        standard_cases[word_lower] = first

    if normalize_plurals:
        for plural, singular in merged_plurals.items():
            standard_cases[plural] = standard_cases.get(singular, singular)

    return fused_cases, standard_cases


def unigrams_and_bigrams(words, stopwords, normalize_plurals=True,
                         collocation_threshold=30):
    """Return word counts including statistically significant bigrams."""
    bigrams = [
        p for p in _pairwise(words)
        if not any(w.lower() in stopwords for w in p)
    ]
    unigrams = [w for w in words if w.lower() not in stopwords]
    n_words = len(unigrams)

    counts_unigrams, standard_form = process_tokens(
        unigrams, normalize_plurals=normalize_plurals)
    counts_bigrams, _ = process_tokens(
        [" ".join(b) for b in bigrams], normalize_plurals=normalize_plurals)

    orig_counts = counts_unigrams.copy()

    for bigram_string, count in counts_bigrams.items():
        parts = bigram_string.split(" ", 1)
        if len(parts) != 2:
            continue
        word1 = standard_form.get(parts[0].lower(), parts[0])
        word2 = standard_form.get(parts[1].lower(), parts[1])
        if word1 not in orig_counts or word2 not in orig_counts:
            continue
        score = _collocation_score(count, orig_counts[word1],
                                   orig_counts[word2], n_words)
        if score > collocation_threshold:
            counts_unigrams[word1] -= count
            counts_unigrams[word2] -= count
            counts_unigrams[bigram_string] = count

    # Remove non-positive counts
    return {w: c for w, c in counts_unigrams.items() if c > 0}