import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import os

def create_variance_attribution_chart():
    # ==========================================
    # 物理级环境重置：抹除所有全局样式污染
    # ==========================================
    mpl.style.use('default')
    mpl.rcParams.update(mpl.rcParamsDefault)
    mpl.rcParams['axes.grid'] = False
    mpl.rcParams['axes.axisbelow'] = True
    mpl.rcParams['font.family'] = 'sans-serif'
    mpl.rcParams['font.sans-serif'] = ['Helvetica', 'Arial']

    # ==========================================
    # 数据准备
    # ==========================================
    categories = ['Opacity', 'Anisotropy', 'Scale', 'SH Energy Ratio', 'PSNR']
    method_variance = [86, 68, 43, 17, 1.4]
    scene_variance = [10, 20, 45, 57, 98.6]

    categories.reverse()
    method_variance.reverse()
    scene_variance.reverse()

    color_method = '#4C72B0'
    color_scene = '#DD8452'

    # ==========================================
    # 画图配置
    # ==========================================
    fig, ax = plt.subplots(figsize=(8, 4.5))

    y_pos = np.arange(len(categories))
    bar_height = 0.6

    bars_method = ax.barh(y_pos, method_variance, height=bar_height, 
                          color=color_method, edgecolor='none', label=r'Method ($SS_{method}$)', zorder=3)
    bars_scene = ax.barh(y_pos, scene_variance, left=method_variance, height=bar_height, 
                         color=color_scene, edgecolor='none', label=r'Scene ($SS_{scene}$)', zorder=3)

    # ==========================================
    # 细节修饰
    # ==========================================
    ax.set_yticks(y_pos)
    ax.set_yticklabels(categories, fontsize=11)
    ax.set_xlim(0, 100)
    
    ax.set_xticks([0, 20, 40, 60, 80, 100])
    ax.set_xticklabels(['0%', '20%', '40%', '60%', '80%', '100%'], color='#555555')

    for i, (m_val, s_val) in enumerate(zip(method_variance, scene_variance)):
        if m_val > 5:
            ax.text(m_val / 2, i, f'{m_val}%' if m_val == int(m_val) else f'{m_val}%', 
                    ha='center', va='center', color='white', fontweight='bold', fontsize=10, zorder=4)
        elif m_val > 0:
            ax.text(m_val + 2, i, f'{m_val}%', 
                    ha='left', va='center', color=color_method, fontweight='bold', fontsize=10, zorder=4)

        if s_val > 5:
            ax.text(m_val + s_val / 2, i, f'{s_val}%' if s_val == int(s_val) else f'{s_val}%', 
                    ha='center', va='center', color='white', fontweight='bold', fontsize=10, zorder=4)

    # ==========================================
    # 彻底的坐标轴与网格清理机制 (终极防御)
    # ==========================================
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['bottom'].set_linewidth(1.5)
    ax.spines['bottom'].set_color('#333333')
    
    # 关闭 X 轴和 Y 轴的刻度短线
    ax.tick_params(axis='y', length=0)
    ax.tick_params(axis='x', length=0) 
    
    # 强制遍历底层实例并将其不可见化
    ax.grid(False)
    for line in ax.get_xgridlines() + ax.get_ygridlines():
        line.set_visible(False)
        line.set_color('none')
        line.set_linewidth(0)

    # ==========================================
    # 分割线
    # ==========================================
    ax.axhline(y=0.5, color='#CCCCCC', linestyle='--', linewidth=1, zorder=2)

    ax.set_title("Asymmetric Variance Attribution", fontsize=14, fontweight='bold', pad=20)
    ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.08), ncol=2, frameon=False, fontsize=11)

    plt.tight_layout()
    
    out_dir = "/root/autodl-tmp/SplatAtlas/tex/figures"
    os.makedirs(out_dir, exist_ok=True)
    out_pdf = os.path.join(out_dir, "fig2_variance_attribution.pdf")
    out_png = os.path.join(out_dir, "fig2_variance_attribution.png")
    
    plt.savefig(out_pdf, format='pdf', bbox_inches='tight', dpi=300)
    plt.savefig(out_png, format='png', bbox_inches='tight', dpi=300)
    
    print(f"✅ Fig 2 生成成功 (全局样式污染已重置，底层网格实例已物理擦除): {out_pdf}")

if __name__ == "__main__":
    create_variance_attribution_chart()