# 第一步：导入必要的库
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
from sklearn.metrics import accuracy_score, classification_report
import nltk
from nltk.corpus import movie_reviews, stopwords
from nltk.stem import WordNetLemmatizer
import re

# 第二步：加载并准备数据（使用NLTK的电影评论数据集）
nltk.download('movie_reviews')
nltk.download('stopwords')
nltk.download('wordnet')

# 获取文件ID和对应的情感标签
documents = [(movie_reviews.raw(fileid), category)
             for category in movie_reviews.categories()
             for fileid in movie_reviews.fileids(category)]
np.random.shuffle(documents) # 打乱顺序

# 将数据分为文本和标签
texts, labels = zip(*documents)
labels = [1 if label == 'pos' else 0 for label in labels] # 将标签转为数字：1代表正面，0代表负面

print(f"数据集大小: {len(texts)} 条评论")
print(f"示例文本 (前200字符): {texts[0][:200]}...")
print(f"示例标签: {'正面' if labels[0] == 1 else '负面'}")

# 第三步：文本预处理函数
def preprocess_text(text):
    """一个简单的文本预处理函数：分词、转小写、去停用词、词形还原"""
    # 1. 移除非字母字符并转为小写
    text = re.sub(r'[^a-zA-Z\s]', '', text).lower()
    # 2. 分词
    tokens = text.split()
    # 3. 去停用词
    stop_words = set(stopwords.words('english'))
    tokens = [token for token in tokens if token not in stop_words and len(token) > 2] # 同时去除过短的词
    # 4. 词形还原 (例如: “running” -> “run”)
    lemmatizer = WordNetLemmatizer()
    tokens = [lemmatizer.lemmatize(token) for token in tokens]
    return ' '.join(tokens)

# 应用预处理函数
print("正在预处理文本，请稍候...")
processed_texts = [preprocess_text(text) for text in texts]

# 第四步：划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    processed_texts, labels, test_size=0.2, random_state=42
)
print(f"训练集大小: {len(X_train)}, 测试集大小: {len(X_test)}")

# 第五步：构建模型管道
# 使用Pipeline可以方便地将预处理步骤（TfidfVectorizer）和模型（MultinomialNB）串联起来
model = make_pipeline(
    TfidfVectorizer(max_features=10000), # 限制特征数量，防止维度过高
    MultinomialNB()
)

# 第六步：训练模型
print("正在训练模型...")
model.fit(X_train, y_train)
print("模型训练完成！")

# 第七步：在测试集上进行预测与评估
y_pred = model.predict(X_test)

# 计算准确率
accuracy = accuracy_score(y_test, y_pred)
print(f"\n模型在测试集上的准确率: {accuracy:.4f}")

# 打印更详细的分类报告
print("\n详细分类报告:")
print(classification_report(y_test, y_pred, target_names=['负面', '正面']))

# 第八步：对新文本进行预测示例
def predict_sentiment(text, model):
    """使用训练好的模型预测单个文本的情感"""
    processed_text = preprocess_text(text)
    prediction = model.predict([processed_text])[0]
    probability = model.predict_proba([processed_text])[0]
    sentiment = "正面" if prediction == 1 else "负面"
    confidence = max(probability)
    return sentiment, confidence

# 测试新文本
test_texts = [
    "This movie is absolutely fantastic! The acting is superb and the story is captivating.",
    "I really hated this film. It was a complete waste of time and money.",
    "The movie was okay, nothing special but not terrible either."
]

print("\n--- 新文本情感预测 ---")
for text in test_texts:
    sentiment, confidence = predict_sentiment(text, model)
    print(f"文本: {text[:50]}...")
    print(f"预测情感: {sentiment} (置信度: {confidence:.2%})\n")