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_DIR = "/root/autodl-tmp/SplatAtlas/tex/figures"

# =========================================================================
# 辅助函数 (构建右图 6宫格 用)
# =========================================================================
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:
        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)
            geo = data.get("geometric", {})
            return {
                "opacity": geo.get("alpha_mean", 0.0),
                "aniso": geo.get("gamma_median", 0.0),
                "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.04)) # 字号稍微调大一点点
    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))
    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)

    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 - 45 # 留出底栏空间
        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))
        
    if opacity is not None and aniso is not None:
        bar_h = 35
        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

os.makedirs(OUTPUT_DIR, exist_ok=True)

# =========================================================================
# 处理 Fig 1a: 左图 (Rank Instability) 单独输出
# =========================================================================
print("正在处理 Fig 1a (左图)...")
if not os.path.exists(LEFT_PDF_PATH):
    print(f"❌ 找不到左图 (PDF): {LEFT_PDF_PATH}")
else:
    try:
        pages = convert_from_path(LEFT_PDF_PATH, dpi=300)
        img_left = pages[0]
        
        # 为左图单独创建一个更宽的画板
        fig_a, ax_a = plt.subplots(figsize=(12, 6))
        ax_a.imshow(img_left)
        ax_a.axis('off')
        ax_a.set_title("(a) Per-seed PSNR Rank Instability (BONSAI & LEGO)", 
                      fontsize=20, pad=15, loc='center', fontweight='bold')
        
        out_a_pdf = os.path.join(OUTPUT_DIR, "fig1a_rank_instability.pdf")
        out_a_png = os.path.join(OUTPUT_DIR, "fig1a_rank_instability.png")
        
        fig_a.savefig(out_a_pdf, dpi=300, bbox_inches='tight', format='pdf', transparent=True)
        fig_a.savefig(out_a_png, dpi=300, bbox_inches='tight', format='png', transparent=False)
        print(f"✅ Fig 1a 保存成功: {out_a_pdf}")
        plt.close(fig_a)
    except Exception as e:
        print(f"❌ 左图处理失败: {e}")

# =========================================================================
# 处理 Fig 1b: 右图 (6宫格 Render vs Rep) 单独输出
# =========================================================================
print("\n正在处理 Fig 1b (右图 6宫格)...")
images = []
labels = ["GT"] + METHODS

# 1. 提取 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)

# 2. 提取 Methods
for method in METHODS:
    cell_dir = get_valid_dir(method, SCENE)
    if not cell_dir:
        images.append(Image.new("RGBA", images[0].size, (50,50,50,255)))
        continue
    render_dirs = sorted(glob.glob(f"{cell_dir}/renders_test*"))
    if not render_dirs:
         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):
         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)

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

img_right = Image.new("RGBA", (grid_w, grid_h), (255,255,255,0)) # 使用透明底色
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)

# 4. 单独制图保存
fig_b, ax_b = plt.subplots(figsize=(14, 8)) # 适合 3x2 的宽幅比例
ax_b.imshow(img_right)
ax_b.axis('off')
ax_b.set_title(f"(b) Saturated Cluster ({SCENE.upper()}): Render vs. Representation", 
              fontsize=20, pad=15, loc='center', fontweight='bold')

out_b_pdf = os.path.join(OUTPUT_DIR, "fig1b_saturation_renders.pdf")
out_b_png = os.path.join(OUTPUT_DIR, "fig1b_saturation_renders.png")

fig_b.savefig(out_b_pdf, dpi=300, bbox_inches='tight', format='pdf', transparent=True)
fig_b.savefig(out_b_png, dpi=300, bbox_inches='tight', format='png', transparent=False)
print(f"✅ Fig 1b 保存成功: {out_b_pdf}")
plt.close(fig_b)

print("\n=== 图片分离处理完成！===")