import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

from sklearn.model_selection import train_test_split
texts_train, texts_test, y_train, y_test = train_test_split(texts, target, test_size=0.5, random_state=42)

# task: preprocess each comment in train and test

texts_train = [preprocess(text) for text in texts_train]
texts_test = [preprocess(text) for text in texts_test]

assert texts_train[5] ==  'who cares anymore . they attack with impunity .'
assert texts_test[89] == 'hey todds ! quick q ? why are you so gay'
assert len(texts_test) == len(y_test)

def text_to_bow(text):
    """convert text string to an array of token counts. Use bow_vocabulary."""
    bow = np.zeros(k)

    for token in text.split():
        ind = (
            bow_vocabulary.index(token)
            if token in bow_vocabulary
            else -1
        )
        if ind != -1:
            bow[ind] += 1
    return bow

X_train_bow = np.stack(list(map(text_to_bow, texts_train)))
X_test_bow = np.stack(list(map(text_to_bow, texts_test)))

k_max = len(set(" ".join(texts_train).split()))
assert X_train_bow.shape == (len(texts_train), min(k, k_max))
assert X_test_bow.shape == (len(texts_test), min(k, k_max))
assert np.all(
    X_train_bow[5:10].sum(-1) == np.array([len(s.split()) for s in texts_train[5:10]])
)
assert len(bow_vocabulary) <= min(k, k_max)
assert X_train_bow[6, bow_vocabulary.index(".")] == texts_train[6].split().count(".")

def get_p_x(delta, word_ind, counts, vocab_size, sum_counts):
    return (delta + counts[word_ind]) / (delta * vocab_size + sum_counts)


class BinaryNaiveBayes:
    delta = 1.0  # add this to all word counts to smoothe probabilities

    def fit(self, X, y):
        """
        Fit a NaiveBayes classifier for two classes
        :param X: [batch_size, vocab_size] of bag-of-words features
        :param y: [batch_size] of binary targets {0, 1}
        """
        # first, compute marginal probabilities of every class, p(y=k) for k = 0,1
        p_y_0 = np.count_nonzero(y == 0) / len(y)
        p_y_1 = np.count_nonzero(y == 1) / len(y)
        self.p_y = np.array([p_y_0, p_y_1])

        # count occurences of each word in texts with label 1 and label 0 separately
        word_counts_positive = X[y == 1].sum(axis=0)
        word_counts_negative = X[y == 0].sum(axis=0)
        # ^-- both must be vectors of shape [vocab_size].

        # finally, lets use those counts to estimate p(x | y = k) for k = 0, 1

        # <YOUR CODE HERE>
        sum_word_counts_positive = sum(word_counts_positive)
        sum_word_counts_negative = sum(word_counts_negative)

        self.p_x_given_positive = np.array(
            [
                get_p_x(
                    delta=self.delta,
                    word_ind=word_ind,
                    counts=word_counts_positive,
                    vocab_size=X.shape[1],
                    sum_counts=sum_word_counts_positive,
                )
                for word_ind in range(X.shape[1])
            ]
        )
        self.p_x_given_negative = np.array(
            [
                get_p_x(
                    delta=self.delta,
                    word_ind=word_ind,
                    counts=word_counts_negative,
                    vocab_size=X.shape[1],
                    sum_counts=sum_word_counts_negative,
                )
                for word_ind in range(X.shape[1])
            ]
        )
        # both must be of shape [vocab_size]; and don't forget to add self.delta!

        return self

    def predict_scores(self, X):
        """
        :param X: [batch_size, vocab_size] of bag-of-words features
        :returns: a matrix of scores [batch_size, k] of scores for k-th class
        """
        # compute scores for positive and negative classes separately.
        # these scores should be proportional to log-probabilities of the respective target {0, 1}
        # note: if you apply logarithm to p_x_given_*, the total log-probability can be written
        # as a dot-product with X
        score_negative = np.log(self.p_y[0]) + np.dot(X, np.log(self.p_x_given_negative))
        # <YOUR CODE HERE - compute unnormalized negative log-probability>
        score_positive = np.log(self.p_y[1]) + np.dot(X, np.log(self.p_x_given_positive))
        # <YOUR CODE HERE - compute unnormalized positive log-probability>
        # you can compute total p(x | y=k) with a dot product
        return np.stack([score_negative, score_positive], axis=-1)

    def predict(self, X):
        return self.predict_scores(X).argmax(axis=-1)

naive_model = BinaryNaiveBayes().fit(X_train_bow, y_train)

assert naive_model.p_y.shape == (2,) and naive_model.p_y.sum() == 1 and naive_model.p_y[0] > naive_model.p_y[1]
assert naive_model.p_x_given_positive.shape == naive_model.p_x_given_negative.shape == X_train_bow.shape[1:]
assert np.allclose(naive_model.p_x_given_positive.sum(), 1.0)
assert np.allclose(naive_model.p_x_given_negative.sum(), 1.0)
assert naive_model.p_x_given_negative.min() > 0, "did you forget to add delta?"

f_index = bow_vocabulary.index('fuck')  # offensive tweets should contain more of this
assert naive_model.p_x_given_positive[f_index] > naive_model.p_x_given_negative[f_index]

g_index = bow_vocabulary.index('good')  # offensive tweets should contain less of this
assert naive_model.p_x_given_positive[g_index] < naive_model.p_x_given_negative[g_index]

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV, StratifiedKFold

log_reg = LogisticRegression(solver='lbfgs', random_state=42)
parameters = {"C": np.logspace(-1, 1, 300)}

bow_log_reg_model = clf.best_estimator_

tokenized_texts_train = []
for text in texts_train:
    tokenized_text = [token for token in text.split() if token in bow_vocabulary]
    tokenized_texts_train.append(tokenized_text)

def get_tf(tokenized_text, term):
    return tokenized_text.count(term) / len(tokenized_text)

def get_idf(tokenized_corpus, term):
    size = len(tokenized_corpus)
    doc_count = sum([1 if term in text else 0 for text in tokenized_corpus])
    
    return np.log((1 + size) / (1 + doc_count)) + 1

def get_tf_idf(tf, idf):
    return tf * idf

def get_l2_norm(vector):
    return np.sqrt(np.sum(vector**2))

def text_to_tf_idf(text):
    """convert text string to an array of token counts. Use bow_vocabulary."""
    vec = np.zeros(k)
    tokenized_text = text.split()
    tf_idf_sentence = []
    tokenized_text_set = list(set(tokenized_text))
    
    for word in tokenized_text_set:
        tf_ = get_tf(tokenized_text, word)
        idf_ = get_idf(tokenized_texts_train, word)
        tf_idf_sentence.append(get_tf_idf(tf_, idf_))

    tf_idf_sentence = np.array(tf_idf_sentence)
    l2_norm = get_l2_norm(tf_idf_sentence)
    tf_idf_sentence = tf_idf_sentence / l2_norm
    tf_idf_sentence_dict = dict(zip(tokenized_text_set, tf_idf_sentence))
    
    for token in tokenized_text_set:
        ind = (
            bow_vocabulary.index(token)
            if token in bow_vocabulary
            else -1
        )
        if ind != -1:
            vec[ind] = tf_idf_sentence_dict[token]
    return vec

X_train_tf_idf = np.stack(list(map(text_to_tf_idf, texts_train)))
X_test_tf_idf = np.stack(list(map(text_to_tf_idf, texts_test)))

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV, StratifiedKFold

log_reg_tf_idf = LogisticRegression(solver='lbfgs', random_state=42)
parameters = {"C": np.logspace(-1, 1, 600)}

tf_idf_log_reg_model = clf_tf_idf.best_estimator_

# !pip install gensim

import gensim.downloader 
embeddings = gensim.downloader.load("fasttext-wiki-news-subwords-300")

# If you're low on RAM or download speed, use "glove-wiki-gigaword-100" instead. Ignore all further asserts.

def vectorize_sum(comment):
    """
    implement a function that converts preprocessed comment to a sum of token vectors
    """
    embedding_dim = embeddings.vectors.shape[1]
    features = np.zeros([embedding_dim], dtype="float32")

    features = np.sum(
        [
            embeddings.get_vector(word) if word in embeddings.key_to_index else features
            for word in comment.split()
        ],
        axis=0,
    ) / len(comment.split())

    return features


assert np.allclose(
    vectorize_sum("who cares anymore . they attack with impunity .")[::70],
    np.array([ 0.00120684,  0.00290737,  0.01539459, -0.0205673 , -0.05153336])
)

X_train_wv = np.stack([vectorize_sum(text) for text in texts_train])
X_test_wv = np.stack([vectorize_sum(text) for text in texts_test])

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV, StratifiedKFold

parameters = {"C": np.logspace(-2, 2, 1000)}

wv_model = clf.best_estimator_