import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont
import os
import glob
import json
import numpy as np
from pdf2image import convert_from_path

# =========================================================================
# 配置项
# =========================================================================
SCENE = "bonsai"
VIEW = "00000.png"
OUTPUTS_BASE = "/root/autodl-tmp/SplatAtlas/outputs"
METHODS = ["vanilla_3dgs", "pgsr", "ges", "opti3dgs", "steepgs"]

LEFT_PDF_PATH  = "/root/autodl-tmp/SplatAtlas/outputs/phase5b/fig1_left_panel.pdf"
OUTPUT_PREFIX = "/root/autodl-tmp/SplatAtlas/tex/figures/fig1_saturation_puzzle"

# =========================================================================
# 辅助函数
# =========================================================================
def get_valid_dir(method, scene):
    for suffix in ["", "_bak"]:
        d = f"{OUTPUTS_BASE}/{method}_{scene}{suffix}"
        if os.path.exists(d) and len(glob.glob(f"{d}/renders_test*")) > 0:
            return d
    return None

def get_metrics(cell_dir):
    try:
        # 尝试读取 metrics json 以获取底栏数据
        metric_files = sorted(glob.glob(f"{cell_dir}/metrics_test_iter*.json"))
        if not metric_files: return None
        with open(metric_files[-1], 'r') as f:
            data = json.load(f)
            # 根据你之前的截图，提取 alpha_mean 和 gamma_median 
            # 如果你已经按我建议改成了倒数，请自行在此处调整计算
            geo = data.get("geometric", {})
            return {
                "opacity": geo.get("alpha_mean", 0.0),
                "aniso": geo.get("gamma_median", 0.0), # 如果是倒数可能是 condition_number 等
                "psnr": data.get("photometric", {}).get("PSNR", 0.0)
            }
    except Exception:
        return None

def add_label(img, title, psnr, opacity=None, aniso=None, is_gt=False):
    """在图片上绘制带有半透明黑底的白字标签 (模拟截图效果)"""
    draw = ImageDraw.Draw(img, 'RGBA')
    w, h = img.size
    
    # 标题 (顶部居中)
    title_text = title.upper()
    font_size = max(16, int(h * 0.03))
    # 简单回退：如果找不到系统字体，使用默认字体。
    # 为了更好效果，这里尽量用比较大的字号，由于默认字体不可缩放，我们可能需要绘制矩形框来显眼
    try:
        font = ImageFont.truetype("DejaVuSans-Bold.ttf", font_size)
    except:
        font = ImageFont.load_default()
    
    # 标题底框
    padding = 5
    title_bbox = draw.textbbox((0, 0), title_text, font=font)
    tw, th = title_bbox[2] - title_bbox[0], title_bbox[3] - title_bbox[1]
    title_x = (w - tw) // 2
    title_y = 10
    draw.rectangle([title_x - padding, title_y - padding, 
                   title_x + tw + padding, title_y + th + padding], 
                   fill=(0, 0, 0, 180))
    # GT 用白色，Method 用绿色
    title_color = (255, 255, 255) if is_gt else (100, 255, 100)
    draw.text((title_x, title_y), title_text, font=font, fill=title_color)

    # 底部 PSNR 标签 (右下)
    if psnr:
        psnr_text = f"PSNR: {psnr:.2f} dB" if not is_gt else "Reference"
        p_bbox = draw.textbbox((0, 0), psnr_text, font=font)
        pw, ph = p_bbox[2] - p_bbox[0], p_bbox[3] - p_bbox[1]
        px = w - pw - 10
        py = h - ph - 40 # 留出底栏空间
        draw.rectangle([px - padding, py - padding, 
                       px + pw + padding, py + ph + padding], 
                       fill=(0, 0, 0, 180))
        draw.text((px, py), psnr_text, font=font, fill=(255,255,255))
        
    # 底部黑条 (Opacity & Aniso)
    if opacity is not None and aniso is not None:
        bar_h = 30
        draw.rectangle([0, h - bar_h, w, h], fill=(30, 30, 30, 255))
        stat_text = f"opacity μ: {opacity:.3f} | anisotropy μ: {aniso:.3f}" if not is_gt else "opacity μ: 1.000 | anisotropy μ: 1.000"
        s_bbox = draw.textbbox((0, 0), stat_text, font=font)
        sw, sh = s_bbox[2] - s_bbox[0], s_bbox[3] - s_bbox[1]
        draw.text(((w-sw)//2, h - bar_h + (bar_h-sh)//2), stat_text, font=font, fill=(200,200,200))
        
    return img


# =========================================================================
# 步骤 1: 构建右侧 6 宫格图片
# =========================================================================
print("正在构建右侧 6 宫格面板...")

images = []
labels = ["GT"] + METHODS

# 提取 GT
gt_dir = get_valid_dir("vanilla_3dgs", SCENE)
if gt_dir:
    gt_path = os.path.join(gt_dir, "gt_test_30000", VIEW)
    if not os.path.exists(gt_path):
        gt_path = os.path.join(gt_dir, "gt_test", VIEW)
    
    if os.path.exists(gt_path):
        img_gt = Image.open(gt_path).convert("RGBA")
        img_gt = add_label(img_gt, "GT", None, opacity=1.0, aniso=1.0, is_gt=True)
        images.append(img_gt)
    else:
        print(f"❌ 找不到 GT 图像: {gt_path}")
        exit(1)
else:
    print("❌ 找不到 vanilla_3dgs 目录以提取 GT")
    exit(1)

# 提取 Methods
for method in METHODS:
    cell_dir = get_valid_dir(method, SCENE)
    if not cell_dir:
        print(f"⚠️ 找不到方法目录: {method}")
        # 创建占位黑图
        img_placeholder = Image.new("RGBA", images[0].size, (50,50,50,255))
        images.append(img_placeholder)
        continue
        
    render_dirs = sorted(glob.glob(f"{cell_dir}/renders_test*"))
    if not render_dirs:
         print(f"⚠️ 找不到渲染目录: {method}")
         images.append(Image.new("RGBA", images[0].size, (50,50,50,255)))
         continue
         
    img_path = os.path.join(render_dirs[-1], VIEW)
    if not os.path.exists(img_path):
         print(f"⚠️ 找不到图像文件: {img_path}")
         images.append(Image.new("RGBA", images[0].size, (50,50,50,255)))
         continue
         
    img_m = Image.open(img_path).convert("RGBA")
    
    # 提取底栏数据
    metrics = get_metrics(cell_dir)
    if metrics:
        img_m = add_label(img_m, method, metrics['psnr'], metrics['opacity'], metrics['aniso'])
    else:
        img_m = add_label(img_m, method, 0.0)
        
    images.append(img_m)

# 拼接 6 宫格
iw, ih = images[0].size
padding = 10
grid_w = iw * 3 + padding * 2
grid_h = ih * 2 + padding

img_right = Image.new("RGBA", (grid_w, grid_h), (30,30,30,255))
positions = [
    (0, 0), (iw + padding, 0), ((iw + padding) * 2, 0),
    (0, ih + padding), (iw + padding, ih + padding), ((iw + padding) * 2, ih + padding)
]

for img, pos in zip(images, positions):
    img_right.paste(img, pos)


# =========================================================================
# 步骤 2: 与左侧 PDF 拼接
# =========================================================================
print("正在从 PDF 提取左侧面板 (DPI=300)...")
if not os.path.exists(LEFT_PDF_PATH):
    print(f"❌ 找不到左图 (PDF): {LEFT_PDF_PATH}")
    exit(1)
    
try:
    pages = convert_from_path(LEFT_PDF_PATH, dpi=300)
    img_left = pages[0]
except Exception as e:
    print(f"❌ PDF 转换失败: {e}\n请确保已安装 poppler-utils")
    exit(1)

print("正在拼接最终图像...")
target_height = 1500
aspect_l = img_left.width / img_left.height
aspect_r = img_right.width / img_right.height

new_w_l = int(target_height * aspect_l)
new_w_r = int(target_height * aspect_r)

img_left = img_left.resize((new_w_l, target_height), Image.Resampling.LANCZOS)
img_right = img_right.resize((new_w_r, target_height), Image.Resampling.LANCZOS)

fig, (ax1, ax2) = plt.subplots(
    1, 2, 
    figsize=(18, 6), 
    gridspec_kw={'width_ratios': [aspect_l, aspect_r]}
)

ax1.imshow(img_left)
ax1.axis('off')
ax1.set_title("(a) Per-seed PSNR Rank Instability (BONSAI & LEGO)", 
              fontsize=18, pad=12, loc='left', fontweight='bold')

ax2.imshow(img_right)
ax2.axis('off')
ax2.set_title(f"(b) Saturated Cluster ({SCENE.upper()}): Render vs. Representation", 
              fontsize=18, pad=12, loc='left', fontweight='bold')

plt.subplots_adjust(wspace=0.03)

os.makedirs(os.path.dirname(OUTPUT_PREFIX), exist_ok=True)
pdf_path = f"{OUTPUT_PREFIX}.pdf"
png_path = f"{OUTPUT_PREFIX}.png"

print("正在保存最终拼接文件...")
plt.savefig(pdf_path, dpi=300, bbox_inches='tight', format='pdf', transparent=True)
plt.savefig(png_path, dpi=300, bbox_inches='tight', format='png', transparent=False)

print(f"=== Figure 1 (Saturation Puzzle) 拼图合成完毕 ===")
print(f"✅ Saved PDF: {pdf_path} (供 LaTeX 插入)")
print(f"✅ Saved PNG: {png_path} (供快速预览)")