import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report

import seaborn as sns
import matplotlib.pyplot as plt

# <<<CELL>>>

df = pd.read_csv("spam.csv", encoding="latin-1")[['v1','v2']]
df.columns = ['Label', 'Message']
df['Label'] = df['Label'].map({'ham':0,'spam':1})
print(df.head())

# <<<CELL>>>

X = df['Message']
y = df['Label']

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# <<<CELL>>>

vectorizer = CountVectorizer()

# <<<CELL>>>

X_train_vec = vectorizer.fit_transform(X_train)
X_test_vec = vectorizer.transform(X_test)

# <<<CELL>>>

model = MultinomialNB()
model.fit(X_train_vec, y_train)

# <<<CELL>>>

y_pred = model.predict(X_test_vec)

# <<<CELL>>>

print("Accuracy:", accuracy_score(y_test, y_pred))
print(classification_report(y_test, y_pred))

# <<<CELL>>>

cm = confusion_matrix(y_test, y_pred)

# <<<CELL>>>

sns.heatmap(
    cm,
    annot=True,
    fmt='d',
    cmap='Blues',
    xticklabels=['Not Spam', 'Spam'],
    yticklabels=['Not Spam', 'Spam']
)

plt.xlabel("Predicted")
plt.ylabel("Actual")
plt.title("Confusion Matrix")
plt.show()

# <<<CELL>>>

message = input("Enter a message: ")
message_vector = vectorizer.transform([message])
prediction = model.predict(message_vector)
if prediction[0] == 1:
    print("Spam")
else:
    print("Not Spam")
