from collections import Counter
import swifter
import pandas as pd
from scipy import sparse
from scipy.stats import norm
from sklearn.feature_extraction.text import CountVectorizer


class Tfbns:

    """
    This a class for Term-Frequency Binomial Seperation feature scaling

    Attributes:
        laplace_smoothing_value(float): Laplacian smoothing value to avoid divide by zero exception.
        vectorizer_params(dict): Parameters for Scikit-learn's CountVectorizer class. 
        If no parameters are given, the default values are used.
        bns_scores(dict): Binomial seperation scores for each word in the corpus.
        bns_scores_list(list): BNS Scores in a list-like format.
        x_train(array): Training data.
        y_train(array): Training label containing 0 and 1.
        freq(object): Vectorizer object generated by the CountVectorizer class.
    """

    def __init__(self, laplace_smoothing_value=0.001, vectorizer_params={}):

        """
        Constructor for the class.

        Parameters:
            laplace_smoothing_value(float): Laplacian smoothing value to avoid 
            divide by zero exception.
            vectorizer_params(dict): Parameters for Scikit-learn's 
            CountVectorizer class.
        """

        self.laplace_smoothing_value = laplace_smoothing_value
        self.vectorizer_params = vectorizer_params
        self.bns_scores_list = []
        self.bns_scores = None
        self.x_train = None
        self.y_train = None
        self.freq = None

    def __bns_parallel(self, train_df):

        """
        Private method to compute binomial seperation scores.

        Parameters:
            train_df(dataframe): Training data in pandas dataframe format.

        """

        tpr_results = Counter()
        fpr_results = Counter()
        tpr_df = train_df[train_df['label'] == 1]
        fpr_df = train_df[train_df['label'] == 0]
        tpr_df['text_train'].swifter.apply(tpr_results.update)
        fpr_df['text_train'].swifter.apply(fpr_results.update)
        tpr_results = pd.DataFrame.from_dict(
            dict(tpr_results), orient='index', columns=['count_tpr'])
        tpr_results.reset_index(inplace=True)
        tpr_results.rename(columns={'index': 'word'}, inplace=True)
        fpr_results = pd.DataFrame.from_dict(
            dict(fpr_results), orient='index', columns=['count_fpr'])
        fpr_results.reset_index(inplace=True)
        fpr_results.rename(columns={'index': 'word'}, inplace=True)
        pr_results = pd.merge(tpr_results, fpr_results, on='word', how='outer')
        pr_results['count_tpr'].fillna(0.001, inplace=True)
        pr_results['count_fpr'].fillna(0.001, inplace=True)
        pr_results['count_tpr'] = pr_results['count_tpr'] / len(tpr_results)
        pr_results['count_fpr'] = pr_results['count_fpr'] / len(fpr_results)
        pr_results['count_tpr'] = pr_results['count_tpr'].swifter.apply(
            norm.ppf)
        pr_results['count_fpr'] = pr_results['count_fpr'].swifter.apply(
            norm.ppf)
        pr_results['result'] = abs(
            pr_results['count_tpr'] - pr_results['count_fpr'])
        pr_results = pr_results[['word', 'result']]
        return pr_results

    def __binomial_seperation(self):

        """Private method to compute binomial seperation scores."""
        bns_scores = None
        train_x_df = pd.DataFrame(self.x_train, columns=['text'])
        train_x = train_x_df['text'].str.split().values
        train_df = pd.DataFrame({'text_train': train_x, 'label': self.y_train}, columns=[
                                'text_train', 'label'])
        bns_scores = self.__bns_parallel(train_df)
        return bns_scores

    def __term_frequency(self):

        """Private method to compute term frequency scores."""
        self.freq = CountVectorizer(**self.vectorizer_params)
        self.freq.fit(self.x_train.ravel())

    def __generate_bns_scores(self, bns_scores):
        """
        Private method to generate binomial seperation scores.
        Parameters:
            bns_scores(dataframe): Binomial Sepeartion scores in dataframe format.
         """
        vocab = self.freq.vocabulary_
        vocab = {k: v for k, v in sorted(
            vocab.items(), key=lambda item: item[1])}
        for word, index in vocab.items():
            self.bns_scores_list.append(bns_scores.get(word, 1))

        self.bns_scores_diag = sparse.diags(self.bns_scores_list)

    def __tfbns(self):
        """Private method to compute TF-BNS Matrix."""
        self.__term_frequency()
        bns_scores = self.__binomial_seperation()
        bns_scores = bns_scores.set_index('word')['result'].to_dict()
        self.__generate_bns_scores(bns_scores)

    def fit(self, x_train, y_train):
        """
        Method for fit.
        Parameters:
            x_train(array): Training data.
            y_train(array): Training label containing 0 and 1.
         """
        self.x_train = x_train
        self.y_train = y_train
        self.__tfbns()
        return self

    def transform(self, x):
        """
        Method for transform.
        Parameters:
            x: Training or testing data.
         """
        x_freq = self.freq.transform(x)
        x_tfbns = x_freq.dot(self.bns_scores_diag)
        return x_tfbns

    def fit_transform(self, x_train, y_train):
        """
        Calls both fit and transform methods.
        Parameters:
            x_train(array): Training data.
            y_train(array): Training label containing 0 and 1.
         """
        self.fit(x_train, y_train)
        return self.transform(x_train)
