import os
import gc
import polars as pl
import polars.selectors as cs
import pandas as pd
import matplotlib.pyplot as plt
import shap
import xgboost as xgb
import seaborn as sns
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
from sklearn.preprocessing import StandardScaler
from tqdm import tqdm
import numpy as np
import sys
import traceback
import matplotlib

# ================= 1. 路径与配置 =================
csv_path = r"D:\tqq\csv\year_csv\2022.csv"
output_dir = r"D:\tqq\csv\year_csv\result\2022"
os.makedirs(output_dir, exist_ok=True)

# 绘图全局设置
plt.rcParams.update({
    'font.family': 'sans-serif',
    'font.size': 11,
    'axes.labelsize': 13,
    'figure.dpi': 300
})
matplotlib.use('Agg')

# ================= 2. Polars 高效读取与清洗 =================
print("📥 正在读取超大规模数据 (1.66亿行)...")
try:
    all_feature_cols = ['clcd', 'tmp', 'pre', 'dem', 'slope', 'aspect', 'gdp', 'hfp', 'pop', 'vpd', 'SMrz']
    target_col = 'NPP'
    id_col = 'pointid'

    df_pl = pl.scan_csv(csv_path, null_values="-9999", infer_schema_length=10000)

    if "grid_code" in df_pl.collect_schema().names():
        df_pl = df_pl.rename({"grid_code": "clcd"})

    df_pl = df_pl.filter(
        (pl.col("NPP") > 0) &
        (pl.all_horizontal(cs.numeric() > -9000))
    )

    df_pl = df_pl.select([
        pl.col(id_col),
        pl.col('clcd').cast(pl.Int32),
        pl.col('aspect').round().cast(pl.Int32),
        *[pl.col(c).cast(pl.Float32) for c in all_feature_cols if c not in ['clcd', 'aspect']],
        pl.col(target_col).cast(pl.Float32)
    ]).drop_nulls()

    df = df_pl.collect().to_pandas()
    print(f"✅ 数据加载完成。有效总样本数: {len(df)}")

except Exception as e:
    traceback.print_exc()
    sys.exit(1)

# ================= 3. 数据预处理 =================
print("🔄 正在进行数据预处理...")

df['clcd'] = df['clcd'].astype('category')
df['aspect'] = df['aspect'].astype('category')

num_features = [c for c in all_feature_cols if c not in ['clcd', 'aspect']]
scaler = StandardScaler()
df[num_features] = scaler.fit_transform(df[num_features]).astype(np.float32)

X = df[all_feature_cols]
y = df[target_col]
ids = df[id_col]

X_train, X_test, y_train, y_test, ids_train, ids_test = train_test_split(
    X, y, ids, test_size=0.2, random_state=42
)

# ================= 4. 极速参数优化 (GridSearchCV) =================
print("🔎 正在进行极速参数寻优 (GridSearch)...")

if len(X_train) > 1000000:
    X_opt_sample, _, y_opt_sample, _ = train_test_split(
        X_train, y_train, train_size=1000000, random_state=42
    )
else:
    X_opt_sample, y_opt_sample = X_train, y_train

param_grid = {
    'n_estimators': [400, 600, 800],
    'max_depth': [6, 8, 10],
    'learning_rate': [0.01, 0.05, 0.1],
}

base_model = xgb.XGBRegressor(
    tree_method='hist',
    device='cuda',
    enable_categorical=True,
    random_state=42,
    subsample=0.8,
    colsample_bytree=0.8,
    reg_lambda=5,
    min_child_weight=5
)

# 检查是否已有最佳参数文件，如果有则跳过寻优（可选，这里默认每次都跑，因为很快）
grid_search = GridSearchCV(
    estimator=base_model,
    param_grid=param_grid,
    scoring='neg_mean_squared_error',
    cv=3,
    verbose=1
)

grid_search.fit(X_opt_sample, y_opt_sample)
best_params = grid_search.best_params_
print(f"🏆 最佳参数: {best_params}")

with open(os.path.join(output_dir, "best_params.txt"), "w") as f:
    f.write(str(best_params))

# ================= 5. 全量训练 =================
print("🚀 使用最优参数进行全量模型训练...")

final_params = {
    **best_params,
    'subsample': 0.8,
    'colsample_bytree': 0.8,
    'reg_lambda': 5,
    'min_child_weight': 5
}

final_model = xgb.XGBRegressor(
    **final_params,
    tree_method='hist',
    device='cuda',
    enable_categorical=True,
    random_state=42
)

eval_size = min(200000, len(X_test))

final_model.fit(
    X_train, y_train,
    eval_set=[(X_test.sample(n=eval_size), y_test.sample(n=eval_size))],
    verbose=False
)

# ================= 6. 模型评估 =================
print("📊 正在评估模型性能...")
y_pred = final_model.predict(X_test)
r2 = r2_score(y_test, y_pred)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))
mae = mean_absolute_error(y_test, y_pred)
print(f"   R²: {r2:.4f}, RMSE: {rmse:.4f}, MAE: {mae:.4f}")

pd.DataFrame({'Metric': ['R2', 'RMSE', 'MAE'], 'Value': [r2, rmse, mae]}).to_csv(
    os.path.join(output_dir, "model_metrics.csv"), index=False)

# Hexbin 密度图
plt.figure(figsize=(9, 7))
plot_indices = np.random.choice(len(y_test), min(500000, len(y_test)), replace=False)
hb = plt.hexbin(y_test.iloc[plot_indices], y_pred[plot_indices], gridsize=80, cmap='Spectral_r', bins='log', mincnt=1)
plt.plot([y.min(), y.max()], [y.min(), y.max()], 'r--', lw=1.5)
plt.colorbar(hb, label='log10(count)')
stats_text = f"$R^2 = {r2:.3f}$\n$RMSE = {rmse:.2f}$\n$MAE = {mae:.2f}$"
plt.gca().text(0.05, 0.95, stats_text, transform=plt.gca().transAxes, fontsize=12, verticalalignment='top',
               bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
plt.title("Global Model Performance")
plt.tight_layout()
plt.savefig(os.path.join(output_dir, "1_Scientific_Hexbin_Plot.png"))
plt.close('all')

# ================= 7. SHAP 分析 (含断点续传) =================
shap_sample_size = min(200000, len(X_test))
print(f"🧠 计算 SHAP 值 (N={shap_sample_size})...")

# 1. 准备数据
X_shap = X_test.sample(n=shap_sample_size, random_state=42)
ids_shap = ids_test.loc[X_shap.index]
y_shap_obs = y_test.loc[X_shap.index]

# 2. 保存原始数据表 (含 pointid)
print("💾 正在保存 SHAP 样本原始数据（含 pointid）...")
df_shap_export = pd.concat([ids_shap, X_shap, y_shap_obs], axis=1)
df_shap_export.to_csv(os.path.join(output_dir, "SHAP_samples_with_pointid.csv"), index=False)

# 保存纯特征样本 (用于shap plots)
X_shap.to_csv(os.path.join(output_dir, "SHAP_X_sample.csv"), index=False)

# 3. 计算 SHAP 值 (支持加载缓存)
shap_values_path = os.path.join(output_dir, "shap_values.npy")

if os.path.exists(shap_values_path):
    print(f"📂 检测到已存在的 SHAP 文件: {shap_values_path}")
    print("   正在加载，跳过重复计算...")
    shap_values = np.load(shap_values_path)
else:
    print("⚡ 尝试使用 GPU 加速计算 SHAP...")
    # 确保模型在 GPU 上
    final_model.set_params(device="cuda")
    explainer = shap.TreeExplainer(final_model)
    shap_values_list = []

    # 分批处理以防爆显存
    batch_size = 10000
    num_batches = int(np.ceil(len(X_shap) / batch_size))

    try:
        for i in tqdm(range(num_batches), desc="SHAP GPU Calculation"):
            start_idx = i * batch_size
            end_idx = min((i + 1) * batch_size, len(X_shap))
            X_batch = X_shap.iloc[start_idx:end_idx]

            batch_shap = explainer.shap_values(X_batch, check_additivity=False)
            shap_values_list.append(batch_shap)

        shap_values = np.concatenate(shap_values_list, axis=0)
        print("✅ SHAP 计算完成 (GPU模式)")

    except Exception as e:
        print(f"⚠️ GPU 计算失败 ({e})，切换回 CPU 模式...")
        final_model.set_params(device="cpu")
        explainer = shap.TreeExplainer(final_model)
        shap_values = explainer.shap_values(X_shap, check_additivity=False)

    # 【新增】保存 SHAP 值到文件，方便下次直接加载
    print(f"💾 正在保存 SHAP 值到 {shap_values_path} ...")
    np.save(shap_values_path, shap_values)

# ================= 8. SHAP 可视化 (全量数据版) =================
print("🎨 正在生成 SHAP 深度分析图表 (使用全量样本，速度可能较慢)...")

# 8.1 SHAP Summary (全局重要性概览)
plt.figure(figsize=(12, 10))
shap.summary_plot(shap_values, X_shap, show=False)
plt.savefig(os.path.join(output_dir, "2_SHAP_Summary.png"))
plt.close('all')

# 8.2 特征重要性条形图
fi = pd.DataFrame({'feature': all_feature_cols, 'importance': final_model.feature_importances_})
fi = fi.sort_values('importance', ascending=False)
plt.figure(figsize=(10, 6))
sns.barplot(x='importance', y='feature', data=fi, hue='feature', palette='magma', legend=False)
plt.title("Global Feature Importance")
plt.tight_layout()
plt.savefig(os.path.join(output_dir, "3_Feature_Importance.png"))
plt.close('all')

# ------------------------------------------------------------
# 【修改点】直接使用全量数据，不再进行抽样
# ------------------------------------------------------------
print(f"🔥 正在使用全量样本进行绘图: {len(X_shap)} 条数据")
X_shap_plot = X_shap  # 直接引用全量
shap_values_plot = shap_values  # 直接引用全量
# ------------------------------------------------------------

# 8.3 SHAP 单变量依赖图 (多项式拟合)
print("📈 生成多项式拟合依赖图...")
dep_poly_dir = os.path.join(output_dir, "shap_dependence_poly")
os.makedirs(dep_poly_dir, exist_ok=True)

# 重新定义数值型特征列表（排除分类变量）
num_features_plot = [c for c in all_feature_cols if c not in ['clcd', 'aspect']]

for col in tqdm(num_features_plot, desc="Processing Poly Plots"):
    plt.figure(figsize=(7, 5))
    idx = all_feature_cols.index(col)

    # 转换为数值型以防报错
    x_data = pd.to_numeric(X_shap_plot[col], errors='coerce')

    # 【建议】数据量大时，透明度(alpha)还要再调低，点的大小(s)调小，否则会变成一团黑
    sns.regplot(
        x=x_data,
        y=shap_values_plot[:, idx],
        order=2,
        scatter_kws={'alpha': 0.05, 's': 1, 'color': 'gray'},  # alpha从0.1降到0.05
        line_kws={'color': 'red', 'lw': 1.5}
    )
    plt.title(f"SHAP Poly Fit: {col}")
    plt.tight_layout()
    plt.savefig(os.path.join(dep_poly_dir, f"dep_{col}.png"))
    plt.close()

# 8.4 SHAP 分类特征依赖图
print("📊 生成分类特征依赖图...")
dep_cat_dir = os.path.join(output_dir, "shap_dependence_categorical")
os.makedirs(dep_cat_dir, exist_ok=True)
cat_features = [c for c in ['aspect', 'clcd'] if c in all_feature_cols]

for col in cat_features:
    plt.figure(figsize=(7, 5))
    shap.dependence_plot(col, shap_values_plot, X_shap_plot, interaction_index=None, show=False, alpha=0.5)
    plt.title(f"SHAP Categorical: {col}")
    plt.tight_layout()
    plt.savefig(os.path.join(dep_cat_dir, f"dep_{col}.png"))
    plt.close()

# 8.5 SHAP 交互依赖图
print("🔗 生成交互依赖图...")
dep_inter_dir = os.path.join(output_dir, "shap_dependence_interaction")
os.makedirs(dep_inter_dir, exist_ok=True)

for col in tqdm(all_feature_cols, desc="Processing Interaction Plots"):
    plt.figure(figsize=(7, 5))
    # interaction_index="auto" 自动寻找最强交互
    shap.dependence_plot(col, shap_values_plot, X_shap_plot, interaction_index="auto", show=False, alpha=0.5)
    plt.title(f"SHAP Interaction: {col}")
    plt.tight_layout()
    plt.savefig(os.path.join(dep_inter_dir, f"inter_{col}.png"))
    plt.close()
# ================= 9. 完成 =================
del df, X, y, X_train, X_test, final_model, shap_values, explainer, y_pred, X_shap, ids, ids_train, ids_test
gc.collect()

print(f"\n🚀 极速全局分析任务完成！结果已保存在: {output_dir}")