# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import pandas as pd
import numpy as np
from collections import defaultdict
from sklearn.datasets import fetch_20newsgroups
from sklearn.metrics import confusion_matrix
from tqdm import tqdm
import itertools
import matplotlib.pyplot as plt
import re
# %matplotlib inline

# # Original Naive Bayes code
# ##### Student ID: 2109853M-IM20-0015
# ##### This code includes original Naive Bayes and the-bag-of-word.(Question1)

def preprocess(paragraph):
    result=re.sub('[^a-z\s]+',' ',paragraph,flags=re.IGNORECASE) 
    result=re.sub('(\s+)',' ',result) 
    result=result.lower() 
    result=result.strip()
    return result 


# ## Original Naive Bayes

class NaiveBayes:
    def __init__(self, each_classes):
        self.classes = each_classes  
        self.bow_dicts = np.array([defaultdict(lambda: 0) for idx in range(self.classes.shape[0])])

    def BoW(self, words, index):
        if isinstance(words, np.ndarray): words = words[0]
        for word in words.split():  
            self.bow_dicts[index][word] += 1

    def DealBoW(self):
        # constructing BoW for each category
        for word_idx, word in enumerate(self.classes):
            all_word = self.data[self.labels == word]  

            result = [preprocess(item) for item in all_word]
            result = pd.DataFrame(data=result)

            # now costruct BoW of this particular category
            np.apply_along_axis(self.BoW, 1, result, word_idx)
        category = np.empty(self.classes.shape[0])
        word_count = np.empty(self.classes.shape[0])
        return category, word_count

    def Train(self, dataset, labels):
        self.data = dataset
        self.labels = labels

        if not isinstance(self.data, np.ndarray): self.data = np.array(self.data)
        if not isinstance(self.labels, np.ndarray): self.labels = np.array(self.labels)
        all_words = []
        category, word_count = self.DealBoW()
        for wordidx, word in enumerate(tqdm(self.classes, desc='Calculate probability for each category')):
            # Calculating prior probability p(c) for each category
            category[wordidx] = np.sum(self.labels == word) / float(self.labels.shape[0])

            #count = list(self.bow_dicts[wordidx].values())
            word_count[wordidx] = np.sum(
                np.array(list(self.bow_dicts[wordidx].values()))) + 1  # |v| is remaining to be added

            # get all words of this category
            all_words += self.bow_dicts[wordidx].keys()

        self.vocabulary = np.unique(np.array(all_words))
        self.vocab_size = self.vocabulary.shape[0]

        # computing denominator value
        denoms = np.array(
            [word_count[wordidx] + self.vocab_size + 1 for wordidx, word in enumerate(self.classes)])
        self.word_info = [(self.bow_dicts[wordidx], category[wordidx], denoms[wordidx]) for wordidx, word in
                          enumerate(self.classes)]
        self.word_info = np.array(self.word_info)

    def findprobablity(self, data, prob):
        for category_idx, category in tqdm(enumerate(self.classes), desc='Calculate word probability:'):
            for test_token in data.split():  
                # for each word w [ count(w|c)+1 ] / [ count(c) + |V| + 1 ]
                # get total count 
                test_token_counts = self.word_info[category_idx][0].get(test_token, 0) + 1
                test_token_prob = test_token_counts / float(self.word_info[category_idx][2])
                prob[category_idx] += np.log(test_token_prob)
        return

    def wordprob(self, data):
        prob = np.zeros(self.classes.shape[0])  
        self.findprobablity(data, prob)
        post_prob = np.empty(self.classes.shape[0])
        for categoryidx, category in enumerate(self.classes):
            post_prob[categoryidx] = prob[categoryidx] + np.log(self.word_info[categoryidx][1])
        return post_prob

    def Test(self, test_set):
        predictions = []  
        for data in tqdm(test_set, desc='Test:'):
            result = preprocess(data)
            post_prob = self.wordprob(result)  
            predictions.append(self.classes[np.argmax(post_prob)])

        return np.array(predictions)


# + pycharm={"name": "#%%\n"}
newsgroups_train = fetch_20newsgroups(subset='train')
train_data = newsgroups_train.data  # getting all trainign examples
train_labels = newsgroups_train.target

# + pycharm={"name": "#%%\n"}
len(train_data)

# + pycharm={"name": "#%%\n"}
len(train_labels)

# + pycharm={"name": "#%%\n"}
newsgroups_test = fetch_20newsgroups(subset='test')  # loading test data
test_data = newsgroups_test.data  # get test set examples
test_labels = newsgroups_test.target

# + pycharm={"name": "#%%\n"}
len(test_data)

# + pycharm={"name": "#%%\n"}
len(test_labels)

# + pycharm={"name": "#%%\n"}
initialNB = NaiveBayes(np.unique(train_labels))

# + pycharm={"name": "#%%\n"}
initialNB.Train(train_data, train_labels)

# + pycharm={"name": "#%%\n"}
prediction = initialNB.Test(test_data)

# + pycharm={"name": "#%%\n"}
len(prediction)
# -

# ## Test result(accuracy)

# + pycharm={"name": "#%%\n"}
accuracy = np.sum(prediction == test_labels) / float(test_labels.shape[0])
accuracy


# -

# ## Confusion matrix

# + pycharm={"name": "#%%\n"}
def plot_confusion_matrix(cm, classes,
                          title='Confusion matrix',
                          cmap=plt.cm.PuBu):
    """
    This function prints and plots the confusion matrix.
    Normalization can be applied by setting `normalize=True`.
    """
    cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
    print(cm)
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=90)
    plt.yticks(tick_marks, classes)
    thresh = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(j, i, format(cm[i, j], '.2f'),
                 horizontalalignment="center",
                 color="white" if cm[i, j] > thresh else "black")

    plt.ylabel('Ground Truth')
    plt.xlabel('Prediction')
    plt.tight_layout()


# + pycharm={"name": "#%%\n"}
# Compute confusion matrix
cma = confusion_matrix(test_labels, prediction, labels=None, sample_weight=None)
np.set_printoptions(precision=2)

# Plot non-normalized confusion matrix
plt.figure(figsize=(12, 10), facecolor='w', edgecolor='b')
plot_confusion_matrix(cma, classes=newsgroups_test.target_names,
                      title='Confusion Matrix on Naive Bayes')
plt.savefig('NB.png')

StarcoderdataJupyterScriptsDedupFiltered