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 precision_score, recall_score, f1_score, confusion_matrix

# 1. LOAD DATASET

df = pd.read_csv("synthetic_email_spam_400.csv")

print(df.head())
print("Missing Values:\n", df.isnull().sum())

# 2. CLEAN TEXT AND SELECT FEATURES

df = df.dropna(subset=["message", "label"])

df["message"] = df["message"].str.lower()
df["message"] = df["message"].str.replace(
    r"[^a-z\s]", " ", regex=True
)

X = df["message"]
y = df["label"]

print("Classes:\n", y.value_counts())

# 3. TRAIN-TEST SPLIT

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

# 4. CONVERT TEXT INTO NUMBERS

vectorizer = CountVectorizer()

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

# 5. TRAIN NAIVE BAYES

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

# 6. PREDICT SPAM / HAM

y_pred = model.predict(X_test)

comparison = pd.DataFrame({
    "Actual": y_test.values,
    "Predicted": y_pred
})

print("\nPredictions:\n", comparison.head(10))

# 7. EVALUATE MODEL

print("Precision:", precision_score(y_test, y_pred, pos_label="spam"))
print("Recall:", recall_score(y_test, y_pred, pos_label="spam"))
print("F1 Score:", f1_score(y_test, y_pred, pos_label="spam"))

print("Confusion Matrix:\n",
      confusion_matrix(y_test, y_pred, labels=["ham", "spam"]))
