import os, glob
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont
from plyfile import PlyData

try:
    from pdf2image import convert_from_path
    HAS_PDF2IMAGE = True
except ImportError:
    HAS_PDF2IMAGE = False

# =========================================================================
# Config
# =========================================================================
SCENE          = 'truck'
REF_VIEW       = '00014.png'   # vanilla's view as reference for hash matching
OUTPUTS_BASE   = '/root/autodl-tmp/SplatAtlas/outputs'
ASSETS_DIR     = '/root/autodl-tmp/SplatAtlas/outputs/phase5b/gallery_assets_31'
LEFT_PDF       = '/root/autodl-tmp/SplatAtlas/outputs/phase5b/fig1_left_panel.pdf'
OUTPUT_PREFIX  = '/root/autodl-tmp/SplatAtlas/tex/figures/fig1_saturation_puzzle'

METHODS = ['minisplatting', '3dgsmcmc', 'gof', 'pixelgs', 'vanilla_3dgs']

# =========================================================================
# Helpers
# =========================================================================
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}/gt_test*")) > 0:
            return d
    return None


def find_matching_view(cell_dir, ref_arr):
    """
    Hash-match: find the file in this method's gt_test/ that depicts
    the same viewpoint as ref_arr. Returns (filename, gt_arr) or (None, None).
    """
    gt_dirs = sorted(glob.glob(f"{cell_dir}/gt_test*"))
    if not gt_dirs:
        return None, None

    best_name, best_mse, best_arr = None, float('inf'), None

    for f in sorted(glob.glob(f"{gt_dirs[-1]}/*.png")):
        try:
            cand = np.array(Image.open(f).convert('RGB')).astype(np.float32)

            if cand.shape != ref_arr.shape:
                continue

            diff = np.abs(ref_arr - cand)
            diff[diff > 240] = 0
            diff[np.abs(diff - 128) < 15] = 0

            mse = float(np.mean(diff ** 2))

            if mse < best_mse:
                best_mse, best_name, best_arr = mse, os.path.basename(f), cand

        except Exception:
            continue

    return (best_name, best_arr) if best_name and best_mse < 50.0 else (None, None)


def get_render_and_psnr(cell_dir, view_name, gt_arr):
    if view_name is None or gt_arr is None:
        return None, None

    rdirs = sorted(glob.glob(f"{cell_dir}/renders_test*"))
    if not rdirs:
        return None, None

    rp = f"{rdirs[-1]}/{view_name}"
    if not os.path.exists(rp):
        return None, None

    try:
        r_pil = Image.open(rp).convert('RGB')
        r_arr = np.array(r_pil).astype(np.float32)

        if r_arr.shape != gt_arr.shape:
            return None, None

        mse = float(np.mean((r_arr - gt_arr) ** 2))
        psnr = 100.0 if mse == 0 else 20 * np.log10(255.0 / np.sqrt(mse))

        return r_pil, psnr

    except Exception:
        return None, None


def get_ply_metrics(cell_dir):
    ply_paths = sorted(glob.glob(f"{cell_dir}/point_cloud/iteration_*/point_cloud.ply"))
    if not ply_paths:
        return None

    try:
        v = PlyData.read(ply_paths[-1])['vertex']

        opacities = 1.0 / (1.0 + np.exp(-v['opacity']))

        scales_log = np.vstack([
            v['scale_0'],
            v['scale_1'],
            v['scale_2']
        ])

        scales_phys = np.exp(scales_log)

        log_aniso = np.max(scales_log, axis=0) - np.min(scales_log, axis=0)

        return {
            'alpha_med': float(np.median(opacities)),
            'scale_med': float(np.median(scales_phys)),
            'aniso_med': float(np.exp(np.median(log_aniso))),
            'n_gauss':   int(scales_log.shape[1]),
        }

    except Exception:
        return None


def fmt_aniso(a):
    return f"{a:.1f}" if a < 100 else f"{a:,.0f}"


def safe_load_font(font_name, size):
    try:
        return ImageFont.truetype(font_name, size)
    except Exception:
        return ImageFont.load_default()


# =========================================================================
# Step 1: build right panel 3x2 grid via PIL
# =========================================================================
print("[1/3] Building right panel (6 cells)...")

ref_path = f"{ASSETS_DIR}/gt_{SCENE}_{REF_VIEW}"
if not os.path.exists(ref_path):
    raise SystemExit(f"Missing reference GT: {ref_path}")

ref_pil = Image.open(ref_path).convert('RGB')
ref_arr = np.array(ref_pil).astype(np.float32)

W, H = ref_pil.size

panels = [
    (ref_pil, 'GROUND TRUTH', None, None, True)
]

for m in METHODS:
    cell_dir = get_valid_dir(m, SCENE)

    if cell_dir is None:
        panels.append((Image.new('RGB', (W, H), (50, 50, 50)), m, None, None, False))
        print(f"  [{m}] MISSING_DIR")
        continue

    view_name, gt_arr = find_matching_view(cell_dir, ref_arr)
    img, psnr = get_render_and_psnr(cell_dir, view_name, gt_arr)

    if img is None or img.size != (W, H):
        fallback = f"{ASSETS_DIR}/{m}_{SCENE}_{REF_VIEW}"

        if os.path.exists(fallback):
            img = Image.open(fallback).convert('RGB')

            if img.size == (W, H):
                img_arr = np.array(img).astype(np.float32)
                mse = float(np.mean((img_arr - ref_arr) ** 2))
                psnr = 100.0 if mse == 0 else 20 * np.log10(255.0 / np.sqrt(mse))
            else:
                img = Image.new('RGB', (W, H), (50, 50, 50))
                psnr = None
        else:
            img = Image.new('RGB', (W, H), (50, 50, 50))
            psnr = None

    metrics = get_ply_metrics(cell_dir)
    panels.append((img, m, psnr, metrics, False))

    if psnr is not None:
        print(f"  [{m}] view={view_name} PSNR={psnr:.2f}")
    else:
        print(f"  [{m}] view={view_name}")


# =========================================================================
# Layout
# =========================================================================
gap        = 12
bottom_bar = int(H * 0.11)

font_main = safe_load_font("DejaVuSans-Bold.ttf", int(H * 0.045))
font_metric = safe_load_font("DejaVuSans-Bold.ttf", int(H * 0.030))

canvas_w = W * 3 + gap * 2
canvas_h = (H + bottom_bar) * 2 + gap

right_canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30))
draw = ImageDraw.Draw(right_canvas)

positions = [
    (0, 0), (1, 0), (2, 0),
    (0, 1), (1, 1), (2, 1)
]

for (img, name, psnr, metrics, is_gt), (col, row) in zip(panels, positions):
    x = col * (W + gap)
    y = row * (H + bottom_bar + gap)

    right_canvas.paste(img, (x, y))

    pad = 8

    # top-left method label
    label = name.upper()
    bb = draw.textbbox((0, 0), label, font=font_main)
    lw, lh = bb[2] - bb[0], bb[3] - bb[1]

    draw.rectangle(
        [x, y, x + lw + pad * 2, y + lh + pad * 2],
        fill=(0, 0, 0)
    )

    color = (255, 255, 255) if is_gt else (76, 175, 80)

    draw.text(
        (x + pad, y + pad),
        label,
        fill=color,
        font=font_main
    )

    # bottom-right PSNR, overlaid on the image
    if is_gt:
        psnr_text = "Reference"
    elif psnr is not None:
        psnr_text = f"PSNR: {psnr:.2f} dB"
    else:
        psnr_text = "PSNR: ---"

    bb = draw.textbbox((0, 0), psnr_text, font=font_main)
    pw, ph = bb[2] - bb[0], bb[3] - bb[1]

    draw.rectangle(
        [x + W - pw - pad * 2, y + H - ph - pad * 2, x + W, y + H],
        fill=(0, 0, 0)
    )

    draw.text(
        (x + W - pw - pad, y + H - ph - pad),
        psnr_text,
        fill=(255, 255, 255),
        font=font_main
    )

    # bottom bar metrics
    if metrics is not None:
        line1 = (
            f"Opacity μ: {metrics['alpha_med']:.3f}   |   "
            f"Scale μ: {metrics['scale_med']:.4f}   |   "
            f"N: {metrics['n_gauss'] / 1e6:.2f}M"
        )

        line2 = f"Anisotropy μ: {fmt_aniso(metrics['aniso_med'])}"

        b1 = draw.textbbox((0, 0), line1, font=font_metric)
        b2 = draw.textbbox((0, 0), line2, font=font_metric)

        l1h = b1[3] - b1[1]
        l2h = b2[3] - b2[1]

        total = l1h + l2h + 6

        t1y = y + H + (bottom_bar - total) // 2
        t2y = t1y + l1h + 6

        draw.text(
            (x + (W - (b1[2] - b1[0])) // 2, t1y),
            line1,
            fill=(200, 200, 200),
            font=font_metric
        )

        draw.text(
            (x + (W - (b2[2] - b2[0])) // 2, t2y),
            line2,
            fill=(255, 152, 0),
            font=font_metric
        )

    elif is_gt:
        txt = "(reference)"
        b = draw.textbbox((0, 0), txt, font=font_metric)

        draw.text(
            (
                x + (W - (b[2] - b[0])) // 2,
                y + H + (bottom_bar - (b[3] - b[1])) // 2
            ),
            txt,
            fill=(150, 150, 150),
            font=font_metric
        )


# =========================================================================
# Step 2: load left panel from PDF
# =========================================================================
print("[2/3] Loading left panel from PDF...")

if HAS_PDF2IMAGE and os.path.exists(LEFT_PDF):
    try:
        left_img = convert_from_path(LEFT_PDF, dpi=300)[0].convert('RGB')
    except Exception as e:
        print(f"  WARNING: failed to load PDF with pdf2image: {e}")
        print("  Placeholder used instead.")
        left_img = Image.new('RGB', (1200, 900), (240, 240, 240))
else:
    print(f"  WARNING: pdf2image missing or {LEFT_PDF} absent — placeholder used")
    left_img = Image.new('RGB', (1200, 900), (240, 240, 240))


# =========================================================================
# Step 3: compose with matplotlib
# =========================================================================
print("[3/3] Composing final figure...")

aspect_l = left_img.width / left_img.height
aspect_r = right_canvas.width / right_canvas.height

fig_w = 18.0
fig_h = fig_w * 1.0 / (aspect_l + aspect_r + 0.05)

fig = plt.figure(figsize=(fig_w, fig_h + 0.5))

gs = fig.add_gridspec(
    1,
    2,
    width_ratios=[aspect_l, aspect_r],
    wspace=0.05,
    top=0.90,
    bottom=0.02,
    left=0.01,
    right=0.99
)

ax1 = fig.add_subplot(gs[0])
ax1.imshow(np.array(left_img))
ax1.axis('off')

ax2 = fig.add_subplot(gs[1])
ax2.imshow(np.array(right_canvas))
ax2.axis('off')

ax1.text(
    0.0,
    1.015,
    "(a)  Per-seed PSNR rank instability  ·  Bonsai & Lego",
    transform=ax1.transAxes,
    ha='left',
    va='bottom',
    fontsize=13,
    fontweight='bold'
)

ax2.text(
    0.0,
    1.015,
    f"(b)  Saturated cluster ({SCENE.title()})  ·  render vs. representation",
    transform=ax2.transAxes,
    ha='left',
    va='bottom',
    fontsize=13,
    fontweight='bold'
)

os.makedirs(os.path.dirname(OUTPUT_PREFIX), exist_ok=True)

plt.savefig(
    f"{OUTPUT_PREFIX}.pdf",
    dpi=300,
    bbox_inches='tight',
    format='pdf'
)

plt.savefig(
    f"{OUTPUT_PREFIX}.png",
    dpi=300,
    bbox_inches='tight',
    format='png',
    facecolor='white'
)

plt.close()

print(f"\n✓ Saved: {OUTPUT_PREFIX}.pdf")
print(f"✓ Saved: {OUTPUT_PREFIX}.png")