import pandas as pd
import re
import nltk
import matplotlib.pyplot as plt

from nltk.corpus import stopwords

from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, confusion_matrix

# Download stopwords
nltk.download('stopwords', quiet=True)

# FAST PREPROCESS (NO LEMMATIZATION)
def preprocess(texts):
    stop_words = set(stopwords.words('english'))
   
    cleaned = []
    for text in texts:
        text = re.sub('[^a-zA-Z]', ' ', str(text)).lower()
        words = text.split()
        words = [w for w in words if w not in stop_words]
        cleaned.append(' '.join(words))
    return cleaned


def text_classification(file_path, text_col, label_col):

    data = pd.read_csv(file_path)
    data = data[[text_col, label_col]].dropna()

    # TAKE SMALL SAMPLE (VERY IMPORTANT)
    data = data.sample(10000, random_state=42)

    # Plot
    data[label_col].value_counts().plot(kind='bar')
    plt.show()

    # Preprocess
    data[text_col] = preprocess(data[text_col])

    # Split
    X_train, X_test, y_train, y_test = train_test_split(
        data[text_col], data[label_col], test_size=0.2, random_state=42
    )

    # LIMIT FEATURES
    cv = CountVectorizer(max_features=3000)
    X_train = cv.fit_transform(X_train)
    X_test = cv.transform(X_test)

    # FASTER MODEL
    model = LogisticRegression(max_iter=200)
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)

    print("Accuracy:", accuracy_score(y_test, y_pred))
    print("Confusion Matrix:\n", confusion_matrix(y_test, y_pred))


# RUN
text_classification(r"C:/Users/User/OneDrive/문서/IMDB Dataset.csv", "review", "sentiment")
#pip install pandas nltk scikit-learn matplotlib