# Path: shorttext/utils/gensim_corpora.py
# def tokens_to_fracdict(tokens):
#     """ Return normalized bag-of-words (BOW) vectors.
# 
#     :param tokens: list of tokens.
#     :type tokens: list
#     :return: normalized vectors of counts of tokens as a `dict`
#     :rtype: dict
#     """
#     cntdict = defaultdict(lambda : 0)
#     for token in tokens:
#         cntdict[token] += 1
#     totalcnt = sum(cntdict.values())
#     return {token: float(cnt)/totalcnt for token, cnt in cntdict.items()}


from itertools import product
from scipy.spatial.distance import euclidean
from scipy.sparse import csr_matrix
from scipy.optimize import linprog
from shorttext.utils.gensim_corpora import tokens_to_fracdict
import warnings
import numpy as np




def word_mover_distance_linprog(first_sent_tokens, second_sent_tokens, wvmodel, distancefunc=euclidean):
    """ Compute the Word Mover's distance (WMD) between the two given lists of tokens, and return the LP problem class.

    Using methods of linear programming, supported by PuLP, calculate the WMD between two lists of words. A word-embedding
    model has to be provided. The whole `scipy.optimize.Optimize` object is returned.

    Reference: Matt J. Kusner, Yu Sun, Nicholas I. Kolkin, Kilian Q. Weinberger, "From Word Embeddings to Document Distances," *ICML* (2015).

    :param first_sent_tokens: first list of tokens.
    :param second_sent_tokens: second list of tokens.
    :param wvmodel: word-embedding models.
    :param distancefunc: distance function that takes two numpy ndarray.
    :return: the whole result of the linear programming problem
    :type first_sent_tokens: list
    :type second_sent_tokens: list
    :type wvmodel: gensim.models.keyedvectors.KeyedVectors
    :type distancefunc: function
    :rtype: scipy.optimize.OptimizeResult
    """
    nb_tokens_first_sent = len(first_sent_tokens)
    nb_tokens_second_sent = len(second_sent_tokens)

    all_tokens = list(set(first_sent_tokens+second_sent_tokens))
    wordvecs = {token: wvmodel[token] for token in all_tokens}

    first_sent_buckets = tokens_to_fracdict(first_sent_tokens)

# Path: shorttext/utils/gensim_corpora.py
# def tokens_to_fracdict(tokens):
#     """ Return normalized bag-of-words (BOW) vectors.
# 
#     :param tokens: list of tokens.
#     :type tokens: list
#     :return: normalized vectors of counts of tokens as a `dict`
#     :rtype: dict
#     """
#     cntdict = defaultdict(lambda : 0)
#     for token in tokens:
#         cntdict[token] += 1
#     totalcnt = sum(cntdict.values())
#     return {token: float(cnt)/totalcnt for token, cnt in cntdict.items()}


# Path: shorttext/metrics/wasserstein/wordmoverdist.py
from itertools import product
from scipy.spatial.distance import euclidean
from scipy.sparse import csr_matrix
from scipy.optimize import linprog
from shorttext.utils.gensim_corpora import tokens_to_fracdict
import warnings
import numpy as np





def word_mover_distance_linprog(first_sent_tokens, second_sent_tokens, wvmodel, distancefunc=euclidean):
    """ Compute the Word Mover's distance (WMD) between the two given lists of tokens, and return the LP problem class.

    Using methods of linear programming, supported by PuLP, calculate the WMD between two lists of words. A word-embedding
    model has to be provided. The whole `scipy.optimize.Optimize` object is returned.

    Reference: Matt J. Kusner, Yu Sun, Nicholas I. Kolkin, Kilian Q. Weinberger, "From Word Embeddings to Document Distances," *ICML* (2015).

    :param first_sent_tokens: first list of tokens.
    :param second_sent_tokens: second list of tokens.
    :param wvmodel: word-embedding models.
    :param distancefunc: distance function that takes two numpy ndarray.
    :return: the whole result of the linear programming problem
    :type first_sent_tokens: list
    :type second_sent_tokens: list
    :type wvmodel: gensim.models.keyedvectors.KeyedVectors
    :type distancefunc: function
    :rtype: scipy.optimize.OptimizeResult
    """
    nb_tokens_first_sent = len(first_sent_tokens)
    nb_tokens_second_sent = len(second_sent_tokens)

    all_tokens = list(set(first_sent_tokens+second_sent_tokens))
    wordvecs = {token: wvmodel[token] for token in all_tokens}

    first_sent_buckets = tokens_to_fracdict(first_sent_tokens)

    second_sent_buckets = tokens_to_fracdict(second_sent_tokens)