from sklearn.tree import DecisionTreeRegressor  # 回归树
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor  # 随机森林与梯度提升
from sklearn.model_selection import train_test_split  # 数据集划分
from sklearn.metrics import mean_squared_error  # MSE评估

# 使用更大样本量的数据进行集成学习
modeling_data = financial_data.sample(n=min(2000, len(financial_data)), random_state=42)  # 抽取2000样本
X_all = modeling_data[['Margin', 'Turnover', 'Leverage', 'Log_Assets']]  # 四维特征矩阵
y_all = modeling_data['ROE']  # ROE目标

# 70%/30%划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_all, y_all, test_size=0.3, random_state=42)

# 训练四种模型
single_tree = DecisionTreeRegressor(random_state=42)  # 单棵回归树（无限制深度）
single_tree.fit(X_train, y_train)  # 拟合

bagging = RandomForestRegressor(n_estimators=100, max_features=None, random_state=42)  # Bagging（全特征）
bagging.fit(X_train, y_train)  # 拟合

rf = RandomForestRegressor(n_estimators=100, max_features='sqrt', random_state=42)  # 随机森林（特征子采样）
rf.fit(X_train, y_train)  # 拟合

boosting = GradientBoostingRegressor(n_estimators=100, learning_rate=0.1, max_depth=3, random_state=42)  # 梯度提升
boosting.fit(X_train, y_train)  # 拟合

from sklearn.ensemble import RandomForestClassifier  # 随机森林分类器
from sklearn.metrics import accuracy_score, roc_auc_score  # 准确率与AUC评估
import warnings  # 警告控制
warnings.filterwarnings('ignore')  # 关闭冗余警告

# 加载海康威视前复权股价数据
stock_path = os.path.join(DATA_DIR, 'stock/stock_price_pre_adjusted.h5')  # 股价文件路径
stock_data = pd.read_hdf(stock_path)  # 读取前复权股价数据

if 'order_book_id' in stock_data.index.names:  # 若在索引中则重置
    stock_data = stock_data.reset_index()

stock_data = stock_data[stock_data['order_book_id'] == '002415.XSHE'].copy()  # 筛选海康威视
stock_data = stock_data.sort_values('date')  # 按日期排序

# 构建金融技术特征
stock_data['Ret'] = stock_data['close'].pct_change()  # 日收益率
stock_data['Target'] = (stock_data['Ret'].shift(-1) > 0).astype(int)  # 明天是否上涨

for i in range(1, 4):  # 构造滞后收益率特征
    stock_data[f'Lag_{i}'] = stock_data['Ret'].shift(i)

stock_data['Vol_5'] = stock_data['Ret'].rolling(5).std()  # 5日波动率
stock_data['Vol_Change'] = stock_data['volume'].pct_change()  # 成交量变化率

stock_data = stock_data.replace([np.inf, -np.inf], np.nan).dropna()  # 清洗inf和NaN
feat_cols = [f'Lag_{i}' for i in range(1, 4)] + ['Vol_5', 'Vol_Change']  # 特征列

# 时间前向切分（70%训练/30%测试）严格保持时间顺序
split = int(len(stock_data) * 0.7)  # 70%分界点
X_tr, X_te = stock_data[feat_cols].values[:split], stock_data[feat_cols].values[split:]  # 特征切分
y_tr, y_te = stock_data['Target'].values[:split], stock_data['Target'].values[split:]  # 目标切分

from sklearn.tree import DecisionTreeClassifier, export_text  # 分类树与规则导出
from sklearn.metrics import classification_report, confusion_matrix  # 分类评估

# 准备数据：预测公司是否盈利
lab_data = financial_data.sample(n=min(2000, len(financial_data)), random_state=42)  # 抽取子样本
lab_data['Profitable'] = (lab_data['ROE'] > 0).astype(int)  # 二分类目标：是否盈利

lab_features = lab_data[['Margin', 'Turnover', 'Leverage', 'Log_Assets']]  # 四维特征
lab_target = lab_data['Profitable']  # 目标变量

# 70%/30%随机划分
lab_X_tr, lab_X_te, lab_y_tr, lab_y_te = train_test_split(
    lab_features, lab_target, test_size=0.3, random_state=42)

# 拟合分类树（Gini准则，最大深度3层）
lab_cls_tree = DecisionTreeClassifier(criterion='gini', max_depth=3, random_state=42)
lab_cls_tree.fit(lab_X_tr, lab_y_tr)  # 训练分类树

lab_pred = lab_cls_tree.predict(lab_X_te)  # 测试集预测
print('分类报告:')  # 输出标题
print(classification_report(lab_y_te, lab_pred))  # 输出精确率、召回率、F1
print(f'混淆矩阵:\n{confusion_matrix(lab_y_te, lab_pred)}')  # 输出混淆矩阵

# 导出树规则为文本（便于硬编码进生产系统）
tree_rules = export_text(lab_cls_tree, feature_names=list(lab_features.columns))  # 导出if-else规则
print('树规则:')  # 输出标题
print(tree_rules)  # 打印规则文本

from sklearn.metrics import r2_score  # R²评估

# 回归目标：预测具体ROE数值
lab_reg_target = lab_data['ROE']  # 连续目标变量
lab_reg_X_tr, lab_reg_X_te, lab_reg_y_tr, lab_reg_y_te = train_test_split(
    lab_features, lab_reg_target, test_size=0.3, random_state=42)  # 数据划分

# 拟合回归树（最大深度4层）
lab_reg_tree = DecisionTreeRegressor(max_depth=4, random_state=42)  # 初始化
lab_reg_tree.fit(lab_reg_X_tr, lab_reg_y_tr)  # 训练

reg_pred = lab_reg_tree.predict(lab_reg_X_te)  # 测试集预测
print(f'回归树 测试集MSE: {mean_squared_error(lab_reg_y_te, reg_pred):.4f}')  # 输出MSE
print(f'回归树 测试集R²: {r2_score(lab_reg_y_te, reg_pred):.4f}')  # 输出R²

# 随机森林（100棵树）
lab_rf = RandomForestRegressor(n_estimators=100, random_state=42)  # 初始化随机森林
lab_rf.fit(lab_reg_X_tr, lab_reg_y_tr)  # 训练

rf_pred = lab_rf.predict(lab_reg_X_te)  # 测试集预测
print(f'随机森林 测试集MSE: {mean_squared_error(lab_reg_y_te, rf_pred):.4f}')  # MSE
print(f'随机森林 测试集R²: {r2_score(lab_reg_y_te, rf_pred):.4f}')  # R²

# 特征重要性排名
importance_df = pd.DataFrame(  # 构建重要性DataFrame
    {'importance': lab_rf.feature_importances_},  # 特征重要性值
    index=lab_features.columns  # 特征名称作为索引
).sort_values('importance', ascending=False)  # 按重要性降序排列

print('\n特征重要性 (Random Forest):')  # 输出标题
print(importance_df)  # 打印排名表

# 梯度提升树（100棵浅树，学习率0.1）
lab_gb = GradientBoostingRegressor(n_estimators=100, learning_rate=0.1, max_depth=3, random_state=42)
lab_gb.fit(lab_reg_X_tr, lab_reg_y_tr)  # 训练

gb_pred = lab_gb.predict(lab_reg_X_te)  # 测试集预测
print(f'Boosting 测试集MSE: {mean_squared_error(lab_reg_y_te, gb_pred):.4f}')  # MSE
print(f'Boosting 测试集R²: {r2_score(lab_reg_y_te, gb_pred):.4f}')  # R²

# 三种方法横向对比
print('\n--- 三种方法对比 ---')  # 输出标题
print(f'单棵树  MSE={mean_squared_error(lab_reg_y_te, reg_pred):.4f}, R²={r2_score(lab_reg_y_te, reg_pred):.4f}')
print(f'随机森林 MSE={mean_squared_error(lab_reg_y_te, rf_pred):.4f}, R²={r2_score(lab_reg_y_te, rf_pred):.4f}')
print(f'Boosting MSE={mean_squared_error(lab_reg_y_te, gb_pred):.4f}, R²={r2_score(lab_reg_y_te, gb_pred):.4f}')