import pandas as pd

model_data = pd.read_csv("./clean.csv")
category_dummies = pd.get_dummies(model_data['category'], prefix='category')
model_data = pd.concat([model_data, category_dummies], axis=1)
features = ['user_age', 'user_income', 'browse_time', 'price', 'quantity'] + \
           [col for col in model_data.columns if col.startswith('category_')]
target = 'total_amount'

X = model_data[features]
y = model_data[target]

print("特征纬度：", X.shape)
print("目标变量统计：", y.describe())

# 数据标准化
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)
lr_model = LinearRegression()
lr_model.fit(X_train, y_train)
y_pred_lr = lr_model.predict(X_test)

rf_model = RandomForestRegressor()
rf_model.fit(X_train, y_train)
y_pred_rf = rf_model.predict(X_test)


def evaluate_model(y_true, y_pred, model_name):
    mse = mean_squared_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    print(f"{model_name} 评估结果：")
    print(f"MSE：{mse:.2f} ")
    print(f"R2：{mse:.4f} ")
    return mse, r2


lr_mse, lr_r2 = evaluate_model(y_test, y_pred_lr, "线性回归")
rf_mse, rf_r2 = evaluate_model(y_test, y_pred_rf, "随机森林")

# 交叉验证
cv_stores = cross_val_score(rf_model, X_scaled, y, cv=5, scoring='r2')
print(f"随机森林R2分数：{cv_stores}")
print(f"平均R2分数：{cv_stores.mean():.4f}(+/-{cv_stores.std() * 2:.4f})")

import  seaborn as sns
import  matplotlib.pyplot as plt
import  matplotlib
matplotlib.use("Agg")
feature_importance = pd.DataFrame({
    'feature' :features,
    'importance':rf_model.feature_importances_
}).sort_values('importance',ascending=False)
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # macOS系
plt.figure(figsize=(20,40))
sns.barplot(x='importance',y='feature',data=feature_importance.head(10))
plt.title("随机森林模型特征重要性排名")
plt.xlabel("特征重要性")
plt.tight_layout()
plt.savefig("./随机森林模型特征重要性排名.png")
