import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False

# 读取数据（排除昨天 20260125）
dates = ['20260123', '20260124']
dfs_a, dfs_c = [], []

for dt in dates:
    dfs_a.append(pd.read_csv(f'output/06a_str_pred分桶诊断_full/{dt}.csv'))
    dfs_c.append(pd.read_csv(f'output/06c_ros_pred分桶诊断_full/{dt}.csv'))

df_a = pd.concat(dfs_a, ignore_index=True)
df_c = pd.concat(dfs_c, ignore_index=True)

print(f"数据日期: {dates} (排除昨天)")
print("=" * 80)

exp_groups = ['对照组', '实验组-str+校准', '实验组-str+校准&ros-统计量']

# ========== 1. 真实值 vs 预测值对比（全量） ==========
print("\n【1. 真实值 vs 预测值对比 - 全量汇总】")
print("-" * 80)

df_all = df_a[(df_a['bucket'] == '全部') & (df_a['apptype'] == 4)].copy()
df_summary = df_all.groupby('abcode', group_keys=False).apply(lambda x: pd.Series({
    'sample_cnt': x['sample_cnt'].sum(),
    'str_real': np.average(x['str_real'], weights=x['sample_cnt']),
    'str_pred': np.average(x['str_pred_avg'], weights=x['sample_cnt']),
    'ros_real': np.average(x['ros_real'], weights=x['sample_cnt']),
    'ros_pred': np.average(x['ros_pred_avg'], weights=x['sample_cnt']),
    'ros_stat': np.average(x['ros_stat_avg'], weights=x['sample_cnt']),
    'rovn_real': np.average(x['rovn_real'], weights=x['sample_cnt']),
    'rovn_pred': np.average(x['rovn_pred'], weights=x['sample_cnt']),
})).reset_index()

print(f"\n{'组别':<35} {'str_real':>10} {'str_pred':>10} {'str差异':>10}")
print("-" * 70)
for _, row in df_summary[df_summary['abcode'].isin(exp_groups)].iterrows():
    str_diff = row['str_pred'] - row['str_real']
    print(f"{row['abcode']:<35} {row['str_real']:>10.4%} {row['str_pred']:>10.4%} {str_diff:>+10.4%}")

print(f"\n{'组别':<35} {'ros_real':>10} {'ros_pred':>10} {'ros_stat':>10} {'pred差异':>10} {'stat差异':>10}")
print("-" * 100)
for _, row in df_summary[df_summary['abcode'].isin(exp_groups)].iterrows():
    pred_diff = row['ros_pred'] - row['ros_real']
    stat_diff = row['ros_stat'] - row['ros_real']
    print(f"{row['abcode']:<35} {row['ros_real']:>10.4f} {row['ros_pred']:>10.4f} {row['ros_stat']:>10.4f} {pred_diff:>+10.4f} {stat_diff:>+10.4f}")

# ========== 2. 分桶真实值 vs 预测值详细对比 ==========
print("\n\n【2. 分桶真实值 vs 预测值详细对比 - 对照组】")
print("=" * 80)

for name, df in [('str_pred分桶', df_a), ('ros_pred分桶', df_c)]:
    print(f"\n{name}:")
    print("-" * 90)

    df_ctrl = df[(df['abcode'] == '对照组') & (df['bucket'] != '全部') & (df['apptype'] == 4)].copy()
    df_ctrl['bucket'] = df_ctrl['bucket'].astype(int)

    # 聚合
    agg_cols = {
        'sample_cnt': 'sum',
        'return_cnt': 'sum',
        'str_real': lambda x: np.average(x, weights=df_ctrl.loc[x.index, 'sample_cnt']),
        'str_pred_avg': lambda x: np.average(x, weights=df_ctrl.loc[x.index, 'sample_cnt']),
        'ros_real': lambda x: np.average(x, weights=df_ctrl.loc[x.index, 'return_cnt']) if df_ctrl.loc[x.index, 'return_cnt'].sum() > 0 else np.nan,
        'ros_pred_avg': lambda x: np.average(x, weights=df_ctrl.loc[x.index, 'sample_cnt']),
        'ros_stat_avg': lambda x: np.average(x, weights=df_ctrl.loc[x.index, 'sample_cnt']),
    }
    df_bucket = df_ctrl.groupby('bucket').agg(agg_cols).reset_index()

    print(f"{'桶':>4} {'样本量':>12} {'str_real':>10} {'str_pred':>10} {'str差异':>10} {'ros_real':>10} {'ros_pred':>10} {'ros_stat':>10} {'pred差异':>10}")
    print("-" * 110)
    for _, row in df_bucket.iterrows():
        str_diff = row['str_pred_avg'] - row['str_real']
        ros_pred_diff = row['ros_pred_avg'] - row['ros_real'] if pd.notna(row['ros_real']) else np.nan
        print(f"{int(row['bucket']):>4} {int(row['sample_cnt']):>12,} {row['str_real']:>10.4%} {row['str_pred_avg']:>10.4%} {str_diff:>+10.4%} {row['ros_real']:>10.4f} {row['ros_pred_avg']:>10.4f} {row['ros_stat_avg']:>10.4f} {ros_pred_diff:>+10.4f}")

# ========== 3. 可视化：真实值 vs 预测值趋势 ==========
fig, axes = plt.subplots(2, 3, figsize=(16, 10))
date_label = f"数据: {', '.join(dates)}"
fig.suptitle(f'分桶诊断 - 真实值 vs 预测值对比\n({date_label})', fontsize=14, fontweight='bold')

for row_idx, (name, df, title_prefix) in enumerate([
    ('str_pred分桶', df_a, 'str_pred分桶'),
    ('ros_pred分桶', df_c, 'ros_pred分桶')
]):
    df_ctrl = df[(df['abcode'] == '对照组') & (df['bucket'] != '全部') & (df['apptype'] == 4)].copy()
    df_ctrl['bucket'] = df_ctrl['bucket'].astype(int)
    df_bucket = df_ctrl.groupby('bucket').agg({
        'str_real': 'mean',
        'str_pred_avg': 'mean',
        'ros_real': 'mean',
        'ros_pred_avg': 'mean',
        'ros_stat_avg': 'mean',
        'ros_real_return': 'mean',
        'ros_pred_return': 'mean',
    }).reset_index()

    # STR 对比
    ax = axes[row_idx, 0]
    ax.plot(df_bucket['bucket'], df_bucket['str_real'] * 100, 'o-', label='str_real', color='green', linewidth=2)
    ax.plot(df_bucket['bucket'], df_bucket['str_pred_avg'] * 100, 's--', label='str_pred', color='blue', linewidth=2)
    ax.set_xlabel('分桶')
    ax.set_ylabel('STR (%)')
    ax.set_title(f'{title_prefix} - STR真实vs预测')
    ax.legend()
    ax.set_xticks(range(1, 11))

    # ROS 对比（全量）
    ax = axes[row_idx, 1]
    ax.plot(df_bucket['bucket'], df_bucket['ros_real'], 'o-', label='ros_real', color='green', linewidth=2)
    ax.plot(df_bucket['bucket'], df_bucket['ros_pred_avg'], 's--', label='ros_pred', color='blue', linewidth=2)
    ax.plot(df_bucket['bucket'], df_bucket['ros_stat_avg'], '^--', label='ros_stat', color='orange', linewidth=2)
    ax.set_xlabel('分桶')
    ax.set_ylabel('ROS')
    ax.set_title(f'{title_prefix} - ROS真实vs预测vs统计量')
    ax.legend()
    ax.set_xticks(range(1, 11))

    # ROS 仅回流样本对比
    ax = axes[row_idx, 2]
    ax.plot(df_bucket['bucket'], df_bucket['ros_real_return'], 'o-', label='ros_real(回流)', color='green', linewidth=2)
    ax.plot(df_bucket['bucket'], df_bucket['ros_pred_return'], 's--', label='ros_pred(回流)', color='blue', linewidth=2)
    ax.set_xlabel('分桶')
    ax.set_ylabel('ROS (仅回流样本)')
    ax.set_title(f'{title_prefix} - ROS回流样本真实vs预测')
    ax.legend()
    ax.set_xticks(range(1, 11))

plt.tight_layout(rect=[0, 0, 1, 0.95])
plt.savefig('output/分桶诊断_真实vs预测对比.png', dpi=150, bbox_inches='tight')
plt.close()

print(f"\n\n图表已保存至: output/分桶诊断_真实vs预测对比.png")

# ========== 4. 各组预测误差汇总 ==========
print("\n\n【3. 各组预测误差汇总 - MAE对比】")
print("=" * 80)

df_all = df_a[(df_a['bucket'] == '全部') & (df_a['apptype'] == 4)].copy()
mae_summary = df_all.groupby('abcode', group_keys=False).apply(lambda x: pd.Series({
    'ros_pred_mae_return': np.average(x['ros_pred_mae_return'], weights=x['return_cnt']),
    'ros_stat_mae_return': np.average(x['ros_stat_mae_return'], weights=x['return_cnt']),
    'ros_pred_mae_over': np.average(x['ros_pred_mae_over'].dropna(), weights=x.loc[x['ros_pred_mae_over'].dropna().index, 'return_cnt']) if x['ros_pred_mae_over'].notna().any() else np.nan,
    'ros_pred_mae_under': np.average(x['ros_pred_mae_under'].dropna(), weights=x.loc[x['ros_pred_mae_under'].dropna().index, 'return_cnt']) if x['ros_pred_mae_under'].notna().any() else np.nan,
})).reset_index()

print(f"\n{'组别':<35} {'ros_pred_MAE':>12} {'ros_stat_MAE':>12} {'高估MAE':>10} {'低估MAE':>10}")
print("-" * 85)
for _, row in mae_summary[mae_summary['abcode'].isin(exp_groups)].iterrows():
    print(f"{row['abcode']:<35} {row['ros_pred_mae_return']:>12.4f} {row['ros_stat_mae_return']:>12.4f} {row['ros_pred_mae_over']:>10.4f} {row['ros_pred_mae_under']:>10.4f}")

# ========== 5. 关键发现总结 ==========
print("\n\n" + "=" * 80)
print("【关键发现总结】")
print("=" * 80)

ctrl = df_summary[df_summary['abcode'] == '对照组'].iloc[0]
print(f"\n对照组基准:")
print(f"  STR: 真实={ctrl['str_real']:.4%}, 预测={ctrl['str_pred']:.4%}, 差异={ctrl['str_pred']-ctrl['str_real']:+.4%}")
print(f"  ROS: 真实={ctrl['ros_real']:.4f}, 预测={ctrl['ros_pred']:.4f}, 统计量={ctrl['ros_stat']:.4f}")
print(f"       pred差异={ctrl['ros_pred']-ctrl['ros_real']:+.4f}, stat差异={ctrl['ros_stat']-ctrl['ros_real']:+.4f}")

for gname in ['实验组-str+校准', '实验组-str+校准&ros-统计量']:
    row = df_summary[df_summary['abcode'] == gname].iloc[0]
    print(f"\n{gname}:")
    print(f"  STR: 真实={row['str_real']:.4%}, 预测={row['str_pred']:.4%}, 差异={row['str_pred']-row['str_real']:+.4%}")
    print(f"  ROS: 真实={row['ros_real']:.4f}, 预测={row['ros_pred']:.4f}, 统计量={row['ros_stat']:.4f}")
    print(f"       pred差异={row['ros_pred']-row['ros_real']:+.4f}, stat差异={row['ros_stat']-row['ros_real']:+.4f}")
