#!/usr/bin/env python
"""T&T round-2 probe — pinpoint remaining suspects at pixel level."""
import os, sys, inspect, struct
import numpy as np
from PIL import Image

sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official')

DATASET_ROOT = "/root/autodl-tmp/dataset/tnt"
OUTPUT_ROOT  = "/root/autodl-tmp/SplatAtlas/outputs"
SCENES = [("truck", "PASS"), ("lighthouse", "FAIL")]


def sec(title):
    print("\n" + "=" * 70)
    print(f"  {title}")
    print("=" * 70)


def build_args(source_path, img_dir):
    from scene.dataset_readers import readColmapSceneInfo
    sig = inspect.signature(readColmapSceneInfo)
    args = []
    for i, (k, p) in enumerate(sig.parameters.items()):
        if i == 0:                    args.append(source_path)
        elif i == 1:                  args.append(img_dir)
        elif k == "eval":             args.append(True)
        elif k == "train_test_exp":   args.append(False)
        else:                         args.append(p.default if p.default != inspect.Parameter.empty else "")
    return args


def psnr_pil(a_pil, b_pil):
    """统一 resize 到较小者，RGB 对齐 PSNR。"""
    if a_pil.size != b_pil.size:
        tgt = (min(a_pil.size[0], b_pil.size[0]),
               min(a_pil.size[1], b_pil.size[1]))
        a_pil = a_pil.resize(tgt, Image.LANCZOS)
        b_pil = b_pil.resize(tgt, Image.LANCZOS)
    a = np.asarray(a_pil.convert("RGB"), dtype=np.float32) / 255.0
    b = np.asarray(b_pil.convert("RGB"), dtype=np.float32) / 255.0
    mse = float(((a - b) ** 2).mean())
    return (100.0 if mse == 0 else 10 * np.log10(1.0 / mse)), mse


# ---------- Probe 8: PINHOLE full params (cx, cy) ----------
def probe_camera_params(scene):
    sec(f"PROBE 8 — COLMAP PINHOLE full params [{scene}]")
    binf = os.path.join(DATASET_ROOT, scene, "sparse", "0", "cameras.bin")
    if not os.path.exists(binf):
        print(f"[!] no {binf}")
        return
    MODELS = {0:("SIMPLE_PINHOLE",3), 1:("PINHOLE",4), 2:("SIMPLE_RADIAL",4),
              3:("RADIAL",5), 4:("OPENCV",8), 5:("OPENCV_FISHEYE",8),
              6:("FULL_OPENCV",12), 7:("FOV",5),
              8:("SIMPLE_RADIAL_FISHEYE",4), 9:("RADIAL_FISHEYE",5),
              10:("THIN_PRISM_FISHEYE",12)}
    with open(binf, "rb") as f:
        num = struct.unpack("<Q", f.read(8))[0]
        cam_id   = struct.unpack("<I", f.read(4))[0]
        model_id = struct.unpack("<i", f.read(4))[0]
        width    = struct.unpack("<Q", f.read(8))[0]
        height   = struct.unpack("<Q", f.read(8))[0]
        model_name, n_params = MODELS.get(model_id, (f"UNK({model_id})", 0))
        params = struct.unpack(f"<{n_params}d", f.read(8 * n_params))

    print(f"  model={model_name} size={width}x{height} num_cameras={num}")
    print(f"  raw params: {[f'{p:.3f}' for p in params]}")
    if model_name == "PINHOLE":
        fx, fy, cx, cy = params
        dcx, dcy = cx - width/2, cy - height/2
        off_x, off_y = abs(dcx)/width*100, abs(dcy)/height*100
        print(f"  fx={fx:.2f} fy={fy:.2f}")
        print(f"  cx={cx:.2f} (expected {width/2:.2f}, offset {dcx:+.2f}px = {off_x:.2f}%)")
        print(f"  cy={cy:.2f} (expected {height/2:.2f}, offset {dcy:+.2f}px = {off_y:.2f}%)")
        if off_x > 0.5 or off_y > 0.5:
            print(f"  [!!] PRINCIPAL POINT OFFSET >0.5% — our K矩阵用 w/2,h/2 忽略了这个偏移！")
        else:
            print(f"  [OK] 主点在中心附近 (<0.5%偏移)")


# ---------- Probe 9: Native GT vs Ours GT pixel match ----------
def probe_gt_match(scene):
    sec(f"PROBE 9 — GT content alignment (Native-saved vs images_2/<orig_name>) [{scene}]")
    from scene.dataset_readers import readColmapSceneInfo
    cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}")
    gt_dir = os.path.join(cell, "gt_test_30000")
    if not os.path.isdir(gt_dir):
        gt_dir = os.path.join(cell, "renders_test_30000", "gt")
    if not os.path.isdir(gt_dir):
        print(f"[!] 没找到 Native GT 目录")
        return None
    native_gts = sorted([f for f in os.listdir(gt_dir) if f.lower().endswith(('.png','.jpg','.jpeg'))])

    source_path = os.path.join(DATASET_ROOT, scene)
    img_dir = "images_2" if os.path.exists(os.path.join(source_path, "images_2")) else "images"
    scene_info = readColmapSceneInfo(*build_args(source_path, img_dir))
    our_test = scene_info.test_cameras

    print(f"Native GT dir : {gt_dir}  count={len(native_gts)}")
    print(f"Our test cams : count={len(our_test)}  (first={our_test[0].image_name})")
    if len(native_gts) != len(our_test):
        print(f"[!!] COUNT MISMATCH — 直接说明两组 test 根本不同")
        return our_test, native_gts, gt_dir

    print(f"\n假说 A: Native 按 idx rename GT，内容 == Ours。验证方式：对齐 idx 比较像素。")
    print(f"{'idx':>3} {'native_file':>15} {'ours_name':>15} {'nat_size':>11} {'our_size':>11} {'PSNR':>8} {'MSE':>10}")
    print("-" * 78)
    psnrs = []
    for i in range(len(native_gts)):
        nat_pil = Image.open(os.path.join(gt_dir, native_gts[i]))
        our_path = our_test[i].image_path
        if not os.path.exists(our_path):
            our_path = os.path.join(source_path, img_dir, our_test[i].image_name)
        our_pil = Image.open(our_path)
        p, m = psnr_pil(nat_pil, our_pil)
        psnrs.append(p)
        if i < 5:
            print(f"{i:>3} {native_gts[i]:>15} {our_test[i].image_name:>15} "
                  f"{str(nat_pil.size):>11} {str(our_pil.size):>11} {p:>8.2f} {m:>10.6f}")

    arr = np.array(psnrs)
    print(f"\nAll {len(psnrs)} pairs: mean={arr.mean():.2f}  min={arr.min():.2f}  max={arr.max():.2f}")
    print(f"  >40 dB: {(arr>40).sum()}/{len(arr)}   20-40 dB: {((arr>=20)&(arr<=40)).sum()}/{len(arr)}   <20 dB: {(arr<20).sum()}/{len(arr)}")

    if arr.mean() > 50:
        print("[A 验证通过] GT 内容完全对齐。FAIL 原因在 renderer/pose/intrinsics 层面。")
    elif arr.mean() > 30:
        print("[A 弱成立] GT 内容近似一致 (30-50 dB) — 可能有 downsample 滤波器差异，但不足以解释 -1.24 dB")
    else:
        print("[A 被推翻] Native 和 Ours 拿到的是不同的图 —— 就是 test split 真的不一样。")
    return our_test, native_gts, gt_dir


# ---------- Probe 10: Native baseline 自洽性 sanity ----------
def probe_native_sanity(scene, our_test, native_gts, gt_dir):
    sec(f"PROBE 10 — Native render vs Native GT (reproduces baseline) [{scene}]")
    cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}")
    candidates = [
        os.path.join(cell, "renders_test_30000", "renders"),
        os.path.join(cell, "renders_test_30000"),
        os.path.join(cell, "test", "ours_30000", "renders"),
    ]
    renders_dir = next((c for c in candidates if os.path.isdir(c)
                        and any(f.lower().endswith(('.png','.jpg')) for f in os.listdir(c))), None)
    if renders_dir is None:
        print("[!] 没找到 Native renders 目录，跳过")
        return
    native_renders = sorted([f for f in os.listdir(renders_dir) if f.lower().endswith(('.png','.jpg'))])
    print(f"Native renders: {renders_dir}  count={len(native_renders)}")
    if not native_renders:
        return

    psnrs = []
    for i in range(min(len(native_renders), len(native_gts))):
        r = Image.open(os.path.join(renders_dir, native_renders[i]))
        g = Image.open(os.path.join(gt_dir, native_gts[i]))
        p, _ = psnr_pil(r, g)
        psnrs.append(p)
    arr = np.array(psnrs)
    print(f"Native_render vs Native_GT  mean PSNR = {arr.mean():.4f} dB  "
          f"(stored baseline should match)")
    print(f"  → 这个值就是 metrics_test_iter30000.json 里的 Native baseline。复核一下。")


def main():
    for scene, label in SCENES:
        print(f"\n\n{'#'*70}\n# SCENE: {scene}  [{label}]\n{'#'*70}")
        probe_camera_params(scene)
        result = probe_gt_match(scene)
        if result:
            our_test, native_gts, gt_dir = result
            probe_native_sanity(scene, our_test, native_gts, gt_dir)


if __name__ == "__main__":
    main()