import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

# 设置学术图表样式
plt.rcParams.update({'font.size': 12, 'pdf.fonttype': 42, 'ps.fonttype': 42})
sns.set_theme(style="whitegrid", context="paper")

OUT_DIR = Path("/root/autodl-tmp/SplatAtlas/outputs/surfel_hypothesis_7tnt_all/")
df_method = pd.read_csv(OUT_DIR / "method_mean_7tnt.csv")
df_scene = pd.read_csv(OUT_DIR / "scene_mean_7tnt.csv")

# 过滤掉 V3 为空的行（比如 2DGS）
df_v3 = df_method.dropna(subset=['v3_normal_median_deg', 'center_pca_normal_median_deg']).copy()

# ==========================================
# Figure 1: V3 vs Center-PCA Scatter Plot
# ==========================================
plt.figure(figsize=(6, 5))
ax = sns.scatterplot(
    data=df_v3, x='center_pca_normal_median_deg', y='v3_normal_median_deg',
    hue='group', style='group', s=80, alpha=0.8, palette="deep"
)
# 画 y=x 的对角线
lims = [10, 65]
plt.plot(lims, lims, '--', color='gray', zorder=0)
plt.fill_between(lims, lims, 70, color='red', alpha=0.05, zorder=0) # 上半区标红(V3更差)
plt.fill_between(lims, 0, lims, color='green', alpha=0.05, zorder=0) # 下半区标绿(V3更好)

plt.title("Covariance Orientation vs. Center PCA Normal")
plt.xlabel("Center-PCA Normal Median Error (deg)")
plt.ylabel("V3 Covariance Normal Median Error (deg)")
plt.xlim(10, 65)
plt.ylim(25, 70)
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', title="Method Group")
plt.text(12, 60, "V3 is Worse", color='red', fontsize=12, fontweight='bold', alpha=0.5)
plt.text(50, 30, "V3 is Better", color='green', fontsize=12, fontweight='bold', alpha=0.5)

plt.tight_layout()
plt.savefig(OUT_DIR / "Fig1_V3_vs_PCA_Scatter.pdf", bbox_inches='tight')
print(f"✅ Saved Fig 1 to {OUT_DIR / 'Fig1_V3_vs_PCA_Scatter.pdf'}")

# ==========================================
# Figure 2: Surface Near Ratio by Scene Barplot
# ==========================================
plt.figure(figsize=(7, 4))
df_scene_sorted = df_scene.sort_values(by='surface_near_ratio', ascending=False)
sns.barplot(data=df_scene_sorted, x='scene', y='surface_near_ratio', palette="viridis")
plt.title("Proximity of Gaussian Centers to True Surface Across Scenes")
plt.xlabel("Tanks & Temples Scene")
plt.ylabel("Surface-Near Ratio ($< 2\\tau$)")
plt.axhline(y=df_scene['surface_near_ratio'].mean(), color='r', linestyle='--', label='Global Mean')
plt.legend()

plt.tight_layout()
plt.savefig(OUT_DIR / "Fig2_Surface_Proximity_Bar.pdf")
print(f"✅ Saved Fig 2 to {OUT_DIR / 'Fig2_Surface_Proximity_Bar.pdf'}")