import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
import joblib

# Load dataset
df = pd.read_csv("spam.csv", encoding="latin-1")

# Your dataset uses columns: email, label
df = df[['email', 'label']]
df.columns = ["text", "label"]

# Clean text column
df['text'] = df['text'].astype(str).fillna('').str.strip()

# Clean label column (labels are already 0/1)
df['label'] = pd.to_numeric(df['label'], errors='coerce')
df = df.dropna(subset=['label'])      # remove rows if label still NaN
df['label'] = df['label'].astype(int) # ensure int type

# Train-test split
X_train, X_test, y_train, y_test = train_test_split(
    df['text'], df['label'], test_size=0.2, random_state=42
)

# TF-IDF
tfidf = TfidfVectorizer(stop_words="english")
X_train_tfidf = tfidf.fit_transform(X_train)
X_test_tfidf = tfidf.transform(X_test)

# Train model
model = LogisticRegression()
model.fit(X_train_tfidf, y_train)

# Accuracy
pred = model.predict(X_test_tfidf)
print("Accuracy:", accuracy_score(y_test, pred))

# Save artifacts
joblib.dump(model, "model.pkl")
joblib.dump(tfidf, "vectorizer.pkl")