import os
import re
import numpy as np
import pandas as pd
import jieba
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.decomposition import PCA
from sklearn.metrics import accuracy_score, classification_report, roc_auc_score, confusion_matrix
from sklearn.feature_selection import VarianceThreshold
import xgboost as xgb
from xgboost import callback as xgb_callback
from transformers import BertTokenizer, BertModel
import torch
import warnings
warnings.filterwarnings('ignore')
# 设置中文字体（解决matplotlib中文显示问题）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# 设置设备（优先使用GPU）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备：{device}")

# ====================== 1. 生成示例数据（电商用户购买预测） ======================
def generate_sample_data():
    """生成示例结构化数据和文本数据"""
    # 结构化数据
    structured_data = pd.DataFrame({
        'user_id': range(1000, 1500),
        'age': np.random.randint(18, 60, size=500),
        'gender': np.random.choice(['男', '女'], size=500, p=[0.5, 0.5]),
        'consumption_amount': np.random.normal(1000, 300, size=500).round(2),
        'purchase_frequency': np.random.randint(1, 30, size=500),
        'product_category': np.random.choice(['电子产品', '美妆', '服装', '食品', '家居'], size=500),
        'rating': np.random.uniform(1, 5, size=500).round(1),
        'is_purchase': np.random.choice([0, 1], size=500, p=[0.4, 0.6])
    })
    
    # 文本数据（用户评论 + 商品标题）
    comments = [
        "这款产品性能很好，值得购买",
        "质量一般，价格偏高",
        "使用体验很棒，推荐！",
        "物流很慢，包装破损",
        "性价比超高，会回购",
        "尺寸不合适，材质差",
        "比预期的好，满意",
        "售后服务差，不推荐"
    ]
    
    product_titles = [
        "2024新款华为Mate60 Pro 12GB+512GB",
        "兰蔻持妆粉底液PO-01 30ml",
        "优衣库男士纯棉T恤458762",
        "三只松鼠坚果大礼包1500g",
        "小米空气净化器4 Pro",
        "美的电饭煲5L大容量",
        "苹果AirPods Pro 2代",
        "Nike Air Max 270运动鞋"
    ]
    
    text_data = pd.DataFrame({
        'user_id': range(1000, 1500),
        'user_comment': np.random.choice(comments, size=500),
        'product_title': np.random.choice(product_titles, size=500)
    })
    
    # 合并文本数据（将评论和标题拼接）
    text_data['combined_text'] = text_data['user_comment'] + " " + text_data['product_title']
    
    # 合并结构化数据和文本数据
    data = pd.merge(structured_data, text_data, on='user_id')
    
    return data
# 生成数据
data = generate_sample_data()
print("数据基本信息：")
print(data.info())
print("\n数据前5行：")
print(data.head())

# ====================== 2. 结构化数据预处理 ======================
print("====================== 2. 结构化数据预处理 ======================")
def preprocess_structured_data(data):
    """预处理结构化数据"""
    # 分离特征和标签
    X_structured = data[['age', 'gender', 'consumption_amount', 'purchase_frequency', 'product_category', 'rating']]
    y = data['is_purchase']
    
    # 定义数值特征和类别特征
    numeric_features = ['age', 'consumption_amount', 'purchase_frequency', 'rating']
    categorical_features = ['gender', 'product_category']
    
    # 数值特征处理管道：缺失值填充 + 标准化
    numeric_transformer = Pipeline(steps=[
        ('imputer', SimpleImputer(strategy='median')),  # 中位数填充缺失值
        ('scaler', StandardScaler())  # 标准化
    ])
    
    # 类别特征处理管道：缺失值填充 + One-Hot编码
    categorical_transformer = Pipeline(steps=[
        ('imputer', SimpleImputer(strategy='most_frequent')),  # 众数填充缺失值
        ('onehot', OneHotEncoder(handle_unknown='ignore'))  # One-Hot编码，忽略未知类别
    ])
    
    # 组合预处理管道
    preprocessor = ColumnTransformer(
        transformers=[
            ('num', numeric_transformer, numeric_features),
            ('cat', categorical_transformer, categorical_features)
        ])
    
    # 执行预处理
    X_structured_processed = preprocessor.fit_transform(X_structured)
    # 转换为DataFrame（方便后续拼接）
    # 获取特征名称
    num_feature_names = numeric_features
    cat_feature_names = preprocessor.named_transformers_['cat'].named_steps['onehot'].get_feature_names_out(categorical_features)
    feature_names = list(num_feature_names) + list(cat_feature_names)
    # 根据输出类型选择合适的转换方式
    if hasattr(X_structured_processed, 'toarray'):
        X_structured_df = pd.DataFrame(X_structured_processed.toarray(), columns=feature_names)
    else:
        X_structured_df = pd.DataFrame(X_structured_processed, columns=feature_names)
    
    return X_structured_df, y, preprocessor
# 预处理结构化数据
X_structured, y, structured_preprocessor = preprocess_structured_data(data)
print("\n预处理后的结构化特征：")
print(X_structured.head())
print(f"\n结构化特征形状：{X_structured.shape}")
# ====================== 3. 文本数据预处理 ======================
def preprocess_text(text):
    """预处理单条文本"""
    # 去除特殊字符
    text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s]', '', text)
    # 分词
    words = jieba.lcut(text)
    # 去停用词（简单停用词表）
    stopwords = ['的', '了', '是', '我', '你', '他', '她', '它', '在', '有', '就', '都', '和', '与', '及', 'a', 'an', 'the']
    words = [word for word in words if word not in stopwords and len(word) > 0]
    # 拼接为字符串
    processed_text = ' '.join(words)
    return processed_text
# 预处理所有文本
data['processed_text'] = data['combined_text'].apply(preprocess_text)
print("\n预处理后的文本示例：")
print(data[['combined_text', 'processed_text']].head())

class BertEmbeddingExtractor:
    """BERT Embedding提取器"""
    def __init__(self, model_name='bert-base-chinese'):
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.model = BertModel.from_pretrained(model_name).to(device)
        self.model.eval()  # 设置为评估模式
    
    def get_embedding(self, text, max_length=32):
        """提取单条文本的Embedding"""
        # 编码文本
        inputs = self.tokenizer(
            text,
            max_length=max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        ).to(device)
        
        # 前向传播（不计算梯度）
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        # 获取最后一层的Hidden States
        last_hidden_state = outputs.last_hidden_state  # [1, max_length, 768]
        
        # 均值池化得到句子Embedding
        embedding = torch.mean(last_hidden_state, dim=1).squeeze(0)  # [768]
        
        # L2归一化
        embedding = torch.nn.functional.normalize(embedding, p=2, dim=0)
        
        return embedding.cpu().numpy()
    
    def batch_get_embedding(self, texts, batch_size=32):
        """批量提取Embedding"""
        embeddings = []
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]
            
            # 批量编码
            inputs = self.tokenizer(
                batch_texts,
                max_length=32,
                padding='max_length',
                truncation=True,
                return_tensors='pt'
            ).to(device)
            
            # 前向传播
            with torch.no_grad():
                outputs = self.model(**inputs)
            
            # 均值池化
            batch_embeddings = torch.mean(outputs.last_hidden_state, dim=1)  # [batch_size, 768]
            # L2归一化
            batch_embeddings = torch.nn.functional.normalize(batch_embeddings, p=2, dim=1)
            
            embeddings.extend(batch_embeddings.cpu().numpy())
        
        return np.array(embeddings)
# 初始化Embedding提取器
embedding_extractor = BertEmbeddingExtractor()
# 提取文本Embedding
print("\n开始提取文本Embedding...")
texts = data['processed_text'].tolist()
X_embedding = embedding_extractor.batch_get_embedding(texts)
# Embedding降维（PCA）
print(f"\n原始Embedding形状：{X_embedding.shape}")
pca = PCA(n_components=0.95)  # 保留95%的方差
X_embedding_pca = pca.fit_transform(X_embedding)
print(f"降维后的Embedding形状：{X_embedding_pca.shape}")
# 转换为DataFrame
embedding_feature_names = [f'embedding_{i}' for i in range(X_embedding_pca.shape[1])]
X_embedding_df = pd.DataFrame(X_embedding_pca, columns=embedding_feature_names)
# ====================== 4. 特征融合 ======================
# 拼接结构化特征和Embedding特征
X_combined = pd.concat([X_structured.reset_index(drop=True), X_embedding_df.reset_index(drop=True)], axis=1)
# 特征选择：去除方差过小的特征
variance_selector = VarianceThreshold(threshold=0.01)
X_combined_selected = variance_selector.fit_transform(X_combined)
# 获取选择后的特征名称
selected_feature_indices = variance_selector.get_support(indices=True)
selected_feature_names = [X_combined.columns[i] for i in selected_feature_indices]
print(f"\n融合后的特征形状（原始）：{X_combined.shape}")
print(f"融合后的特征形状（特征选择后）：{X_combined_selected.shape}")
print(f"选择后的特征数量：{len(selected_feature_names)}")
# 转换为DataFrame
X_combined_df = pd.DataFrame(X_combined_selected, columns=selected_feature_names)

# ====================== 1. 划分数据集 ======================
X_train, X_test, y_train, y_test = train_test_split(
    X_combined_df, y, test_size=0.2, random_state=42, stratify=y
)
X_train, X_val, y_train, y_val = train_test_split(
    X_train, y_train, test_size=0.25, random_state=42, stratify=y_train
)
print(f"\n训练集形状：{X_train.shape}")
print(f"验证集形状：{X_val.shape}")
print(f"测试集形状：{X_test.shape}")
# ====================== 2. 定义XGBoost模型 ======================
# 设置参数
xgb_params = {
    'objective': 'binary:logistic',  # 二分类任务
    'max_depth': 5,                  # 树最大深度
    'learning_rate': 0.1,            # 学习率
    'n_estimators': 200,             # 树的数量
    'subsample': 0.8,                # 样本采样比例
    'colsample_bytree': 0.8,         # 特征采样比例
    'reg_alpha': 0.1,                # L1正则化
    'reg_lambda': 1,                 # L2正则化
    'eval_metric': 'auc',            # 评估指标
    'seed': 42,                      # 随机种子
    'verbosity': 1                   # 输出日志级别
}
# 初始化模型
model = xgb.XGBClassifier(**xgb_params)
# ====================== 3. 训练模型 ======================
print("\n开始训练XGBoost模型...")
model.fit(X_train, y_train, verbose=True)
# ====================== 4. 模型评估 ======================
# 预测
y_pred = model.predict(X_test)
y_pred_proba = model.predict_proba(X_test)[:, 1]
# 计算评估指标
accuracy = accuracy_score(y_test, y_pred)
auc = roc_auc_score(y_test, y_pred_proba)
classification_rep = classification_report(y_test, y_pred)
print("\n====================== 模型评估结果 ======================")
print(f"准确率（ACC）：{accuracy:.4f}")
print(f"AUC-ROC：{auc:.4f}")
print("\n分类报告：")
print(classification_rep)
# ====================== 5. 特征重要性分析 ======================
# 获取特征重要性
feature_importance = model.get_booster().get_score(importance_type='weight')
# 转换为DataFrame
feature_importance_df = pd.DataFrame({
    'feature': list(feature_importance.keys()),
    'importance': list(feature_importance.values())
})
# 按重要性排序
feature_importance_df = feature_importance_df.sort_values('importance', ascending=False).head(20)
# 绘制特征重要性图
plt.figure(figsize=(12, 8))
sns.barplot(x='importance', y='feature', data=feature_importance_df)
plt.title('XGBoost特征重要性（Top 20）', fontsize=14)
plt.xlabel('重要性得分', fontsize=12)
plt.ylabel('特征名称', fontsize=12)
plt.tight_layout()
plt.savefig('feature_importance.png', dpi=300)
plt.show()
# ====================== 6. 混淆矩阵 ======================
cm = confusion_matrix(y_test, y_pred)
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
            xticklabels=['未购买', '购买'], 
            yticklabels=['未购买', '购买'])
plt.title('混淆矩阵', fontsize=14)
plt.xlabel('预测标签', fontsize=12)
plt.ylabel('真实标签', fontsize=12)
plt.tight_layout()
plt.savefig('confusion_matrix.png', dpi=300)
plt.show()
# ====================== 7. 对比实验（仅使用结构化特征） ======================
print("\n====================== 对比实验：仅使用结构化特征 ======================")
# 仅使用结构化特征训练XGBoost
X_train_struct = X_train[X_structured.columns] if set(X_structured.columns).issubset(X_train.columns) else X_train.iloc[:, :len(X_structured.columns)]
X_test_struct = X_test[X_structured.columns] if set(X_structured.columns).issubset(X_test.columns) else X_test.iloc[:, :len(X_structured.columns)]
# 训练模型
model_struct = xgb.XGBClassifier(**xgb_params)
model_struct.fit(X_train_struct, y_train, verbose=False)
# 评估
y_pred_struct = model_struct.predict(X_test_struct)
y_pred_proba_struct = model_struct.predict_proba(X_test_struct)[:, 1]
accuracy_struct = accuracy_score(y_test, y_pred_struct)
auc_struct = roc_auc_score(y_test, y_pred_proba_struct)
print(f"仅结构化特征 - 准确率（ACC）：{accuracy_struct:.4f}")
print(f"仅结构化特征 - AUC-ROC：{auc_struct:.4f}")
print(f"\n融合模型 vs 仅结构化特征：")
print(f"准确率提升：{(accuracy - accuracy_struct)*100:.2f}%")
print(f"AUC提升：{(auc - auc_struct)*100:.2f}%")
# ====================== 8. 保存模型 ======================
# 保存XGBoost模型
model.save_model('xgboost_llm_fusion.model')
print("\n模型已保存为：xgboost_llm_fusion.model")
# 保存特征名称
with open('feature_names.txt', 'w', encoding='utf-8') as f:
    for name in selected_feature_names:
        f.write(f"{name}\n")
class XGBoostLLMFusionInference:
    """XGBoost+大模型融合推理类"""
    def __init__(self, model_path, feature_names_path, embedding_extractor, structured_preprocessor):
        # 加载模型
        self.model = xgb.XGBClassifier()
        self.model.load_model(model_path)
        
        # 加载特征名称
        with open(feature_names_path, 'r', encoding='utf-8') as f:
            self.feature_names = [line.strip() for line in f.readlines()]
        
        # Embedding提取器
        self.embedding_extractor = embedding_extractor
        
        # 结构化数据预处理器
        self.structured_preprocessor = structured_preprocessor
        
        # PCA模型（这里简化，实际应保存训练好的PCA模型）
        self.pca = pca
    
    def preprocess_input(self, structured_data, text):
        """预处理输入数据"""
        # 1. 预处理结构化数据
        structured_df = pd.DataFrame([structured_data])
        structured_processed = self.structured_preprocessor.transform(structured_df)
        # 根据输出类型选择合适的转换方式
        if hasattr(structured_processed, 'toarray'):
            structured_df_processed = pd.DataFrame(structured_processed.toarray(),
                                                  columns=self.structured_preprocessor.get_feature_names_out())
        else:
            structured_df_processed = pd.DataFrame(structured_processed,
                                                  columns=self.structured_preprocessor.get_feature_names_out())
        # 2. 预处理文本并提取Embedding
        processed_text = preprocess_text(text)
        embedding = self.embedding_extractor.get_embedding(processed_text)
        embedding_pca = self.pca.transform([embedding])
        # 3. 特征融合
        embedding_df = pd.DataFrame(embedding_pca, columns=[f'embedding_{i}' for i in range(embedding_pca.shape[1])])
        combined_df = pd.concat([structured_df_processed, embedding_df], axis=1)
        # 4. 确保所有需要的特征都存在，缺失的填充为0
        for col in self.feature_names:
            if col not in combined_df.columns:
                combined_df[col] = 0
        # 5. 选择模型训练时的特征
        combined_df = combined_df[self.feature_names]
        return combined_df
    
    def predict(self, structured_data, text):
        """预测"""
        # 预处理输入
        input_data = self.preprocess_input(structured_data, text)
        
        # 预测
        pred_proba = self.model.predict_proba(input_data)[0]
        pred_label = self.model.predict(input_data)[0]
        
        return {
            'pred_label': int(pred_label),
            'pred_prob': float(pred_proba[1]),
            'label_name': '购买' if pred_label == 1 else '未购买'
        }
# 初始化推理类
inference = XGBoostLLMFusionInference(
    model_path='xgboost_llm_fusion.model',
    feature_names_path='feature_names.txt',
    embedding_extractor=embedding_extractor,
    structured_preprocessor=structured_preprocessor
)
# 示例推理
sample_structured_data = {
    'age': 28,
    'gender': '女',
    'consumption_amount': 1500.0,
    'purchase_frequency': 20,
    'product_category': '美妆',
    'rating': 4.8
}
sample_text = "这款粉底液遮瑕效果超棒，持妆久，推荐！ 兰蔻持妆粉底液PO-01 30ml"
# 推理
result = inference.predict(sample_structured_data, sample_text)
print("\n====================== 在线推理示例 ======================")
print(f"输入结构化数据：{sample_structured_data}")
print(f"输入文本：{sample_text}")
print(f"预测结果：{result['label_name']}")
print(f"  - 购买概率：{result['pred_prob']:.4f} ({result['pred_prob']*100:.2f}%)")
print(f"  - 未购买概率：{1-result['pred_prob']:.4f} ({(1-result['pred_prob'])*100:.2f}%)")
# 测试多个样本
print("\n====================== 多样本对比测试 ======================")
test_cases = [
    {
        'structured': {'age': 35, 'gender': '女', 'consumption_amount': 2000.0, 'purchase_frequency': 25, 'product_category': '美妆', 'rating': 4.9},
        'text': "产品非常棒，会继续购买！ 兰蔻持妆粉底液PO-01 30ml"
    },
    {
        'structured': {'age': 22, 'gender': '男', 'consumption_amount': 500.0, 'purchase_frequency': 2, 'product_category': '服装', 'rating': 2.5},
        'text': "质量太差，不建议购买 优衣库男士纯棉T恤458762"
    },
    {
        'structured': {'age': 45, 'gender': '女', 'consumption_amount': 3000.0, 'purchase_frequency': 15, 'product_category': '电子产品', 'rating': 4.5},
        'text': "性能出色，推荐购买 2024新款华为Mate60 Pro 12GB+512GB"
    }
]
for i, test_case in enumerate(test_cases, 1):
    pred = inference.predict(test_case['structured'], test_case['text'])
    print(f"\n样本{i}:")
    print(f"  预测: {pred['label_name']} (购买概率: {pred['pred_prob']:.4f})")
    print(f"  评论: {test_case['text']}")