__author__ = 'thorwhalen'

import numpy as np
import pynball.ch as daf_ch
import aspyre.math as semantics_math

import nltk
import pandas as pd


class TermWeightGetter(object):
    def __init__(self, term_gweight_selector):
        self.term_gweight_selector = term_gweight_selector

    def termcount_to_termweight(self, termstats):
        termstats = termcount_to_termtf(termstats)
        term_gweight = self.term_gweight_selector.get_table(list(termstats.index))
        return times(termstats, term_gweight)


def times(sr1, sr2):
    return (sr1*sr2).dropna()


def cosine(sr1, sr2):
    return (
        (sr1*sr2).dropna().sum()
            / (np.linalg.norm(sr1.values)*np.linalg.norm(sr2.values))
    )


def termcount_to_termtf(termcount):
    """
    returns a term-frequency series given a term-count series
    """
    df = termcount / float(termcount.sum())
    df.columns = ['stat']
    return df


def termdoc_to_termdoc_count(term_doc_df, doc_var=None, term_var='term', count_var='count'):
    # processing input
    term_doc_df, doc_var, term_var = __process_term_doc_var_names__(term_doc_df, doc_var=doc_var, term_var=term_var)
    term_doc_df = term_doc_df.groupby([doc_var, term_var]).count()
    term_doc_df = daf_ch.ch_col_names(term_doc_df, count_var, term_var)
    return term_doc_df


def termdoc_to_doc_counts(term_doc_df, doc_var=None, term_var='term', count_var='count'):
    term_doc_df, doc_var, term_var = __process_term_doc_var_names__(term_doc_df, doc_var=doc_var, term_var=term_var)
    # keep only doc and terms, and one copy of any (doc,term) pair
    term_doc_df = term_doc_df[[doc_var, term_var]].drop_duplicates(cols=[doc_var, term_var]).reset_index(drop=True)
    # group by terms
    term_doc_df = term_doc_df[[term_var]].groupby(term_var).count()
    term_doc_df = daf_ch.ch_col_names(term_doc_df, count_var, term_var)
    return term_doc_df


def termdoc_to_term_idf(term_doc_df, doc_var=None, term_var='term'):
    # processing input
    term_doc_df, doc_var, term_var = __process_term_doc_var_names__(term_doc_df, doc_var=doc_var, term_var=term_var)
    # get the number of docs
    num_of_docs = len(np.unique(term_doc_df[doc_var]))
    # get doc_counts
    term_doc_df = termdoc_to_doc_counts(term_doc_df, doc_var, term_var)
    # # keep only doc and terms, and one copy of any (doc,term) pair
    # term_doc_df = term_doc_df[[doc_var,'term']].drop_duplicates(cols=[doc_var, 'term']).reset_index(drop=True)
    # # group by terms
    # term_doc_df = term_doc_df[['term']].groupby('term').count()
    term_doc_df['term'] = \
        semantics_math.idf_log10(num_of_docs_containing_term=np.array(term_doc_df['term']), num_of_docs=float(num_of_docs))
    return daf_ch.ch_col_names(term_doc_df, ['stat'], [term_var])


def __process_term_doc_var_names__(term_doc_df, doc_var=None, term_var='term'):
    if term_var != 'term':
        term_doc_df = daf_ch.ch_col_names(term_doc_df, [term_var], [term_var])
    cols = term_doc_df.columns
    if doc_var is None:  # try to guess it
        if len(cols) != 2:
            raise ValueError("In order to guess the doc_var, there needs to be only two columns")
        else:
            doc_var = list(set(cols)-set([term_var]))[0]
    return (term_doc_df, doc_var, term_var)


def mk_termCounts(dat, indexColName, strColName, data_folder=''):
    """
    input: data_folder='', dataname, savename='',
    output: string of ascii char correspondents
      (replacing, for example, accentuated letters with non-accentuated versions of the latter)
    """
    from tapyoca import log
    from pynball.get import get_data
    dat = get_data(dat,data_folder)
    log.printProgress("making {} word counts (wc)",strColName)
    sr = to_kw_fd(dat[strColName])
    sr.index = dat.hotel_id.tolist()
    return sr
    # translate fds to series
    #printProgress("translating {} FreqDist to Series",col)
    #sr = fd_to_series(sr)
    #printProgress("saving {} word counts (wc)",col)
    #save_data(sr,savename + col)
    #printProgress('Done!')


def to_kw_tokens(dat, indexColName, strColName, data_folder=''):
    """
    input: daf of strings
    output: series of the tokens of the strings, processed for AdWords keywords
      i.e. string is lower capsed and asciied, and words are [\w&]+
    """
    from tapyoca import log
    import typyst.trans
    from pynball.get import get_data
    dat = get_data(dat, data_folder)
    log.printProgress("making {} tokens",strColName)
    sr = dat[strColName]
    sr.index = dat.hotel_id.tolist()
    # preprocess string
    sr = sr.map(lambda x: x.lower())
    sr = sr.map(lambda x: typyst.trans.toascii(x))
    # tokenize
    sr = sr.map(lambda x:nltk.regexp_tokenize(x,'[\w&]+'))
    # return this
    return sr


def to_kw_fd(s):
    """
    input: string, series (with string columns)
    output: nltk.FreqDist (word count) of the tokens of the string, processed for AdWords keywords
      i.e. string is lower capsed and asciied, and words are [\w&]+
    """
    import typyst.trans
    if isinstance(s,pd.Series):
        # preprocess string
        s = s.map(lambda x: x.lower())
        s = s.map(lambda x: typyst.trans.toascii(x))
        # tokenize
        s = s.map(lambda x:nltk.regexp_tokenize(x,'[\w&]+'))
        # return series of fds
        return s.map(lambda tokens:nltk.FreqDist(tokens))
    elif isinstance(s,str):
        # preprocess string
        s = typyst.trans.toascii(s.lower())
        # tokenize
        tokens = nltk.regexp_tokenize(s,'[\w&]+')
        return nltk.FreqDist(tokens)


def fd_to_series(fd):
    """
    input: nltk.FreqDist, pd.Series (with FreqDist in them), or list of FreqDists
    output: the pd.Series representation of the FreqDist
    """
    if isinstance(fd,nltk.FreqDist):
        return pd.Series(list(fd.values()),list(fd.keys()))
    elif isinstance(fd,pd.Series):
        return fd.map(lambda x:pd.Series(list(x.values()),list(x.keys())))
    elif isinstance(fd,list):
        return list(map(lambda x:pd.Series(list(x.values()),list(x.keys()))))