import pandas as pd
import re
import nltk
from Sastrawi.Stemmer.StemmerFactory import StemmerFactory
from nltk.corpus import stopwords
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
import pickle

# --- Data preprocessing ---
data = pd.read_csv("random_dataset_oversampling.csv")
data.columns = ['content', 'Label']

# Case folding
data['content'] = data['content'].str.lower()

# Cleansing simbol, angka, tanda baca
data['content'] = data['content'].apply(lambda x: re.sub(r'[^a-z\s]', ' ', str(x)))

# Tokenisasi awal
data['tokens'] = data['content'].apply(lambda x: x.split())

# Load kamus normalisasi
kamus = pd.read_csv("kamus_normalize.csv", sep=';')
kamus['BEFORE'] = kamus['BEFORE'].str.lower().str.strip()
kamus['AFTER'] = kamus['AFTER'].str.lower().str.strip()
normalisasi_dict = dict(zip(kamus['BEFORE'], kamus['AFTER']))

# Normalisasi
def normalisasi(tokens):
    return [normalisasi_dict.get(token, token) for token in tokens]

data['normalized'] = data['tokens'].apply(normalisasi)

# Stopword removal
nltk.download('stopwords')
stop_words = set(stopwords.words('indonesian'))
stop_words.update(['yg', 'dg', 'rt', 'dgn', 'ny', 'd', 'klo', 'kalo', 'amp', 'biar', 'bikin', 'bilang'])
data['stopwords_removed'] = data['normalized'].apply(lambda tokens: [t for t in tokens if t not in stop_words])

# Stemming
factory = StemmerFactory()
stemmer = factory.create_stemmer()
data['stemmed'] = data['stopwords_removed'].apply(lambda tokens: [stemmer.stem(t) for t in tokens])

# Final clean text
data['clean_text'] = data['stemmed'].apply(lambda x: ' '.join(x))

# --- TF-IDF Vectorizer tanpa identity ---
vectorizer = TfidfVectorizer(max_features=5000)
X_tfidf = vectorizer.fit_transform(data['clean_text'])

# --- Label Encoder ---
label_encoder = LabelEncoder()
y_encoded = label_encoder.fit_transform(data['Label'])

# --- Train/Test Split ---
x_train, x_test, y_train, y_test = train_test_split(X_tfidf, y_encoded, test_size=0.2, random_state=42)

# --- Train Naive Bayes ---
nb_classifier = MultinomialNB()
nb_classifier.fit(x_train, y_train)

# --- Save artifacts ---
with open("tfidf_vectorizer.pkl", "wb") as f:
    pickle.dump(vectorizer, f)

with open("nb_classifier.pkl", "wb") as f:
    pickle.dump(nb_classifier, f)

with open("label_encoder.pkl", "wb") as f:
    pickle.dump(label_encoder, f)

print("TF-IDF, model, dan LabelEncoder berhasil disimpan.")