#!/usr/bin/env python
"""T&T root-cause probe — read-only, symmetric PASS vs FAIL."""
import os, sys, inspect, struct
import numpy as np
from PIL import Image, ImageOps
from plyfile import PlyData

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):
    """与 render_single.py 完全一致的 args 构造逻辑"""
    from scene.dataset_readers import readColmapSceneInfo
    sig = inspect.signature(readColmapSceneInfo)
    args_list = []
    for i, (k, p) in enumerate(sig.parameters.items()):
        if i == 0:                    args_list.append(source_path)
        elif i == 1:                  args_list.append(img_dir)
        elif k == "eval":             args_list.append(True)
        elif k == "train_test_exp":   args_list.append(False)
        else:                         args_list.append(p.default if p.default != inspect.Parameter.empty else "")
    return sig, args_list


# ---------- Probe 1 ----------
def probe_signature():
    sec("PROBE 1 — readColmapSceneInfo signature & args mapping")
    from scene.dataset_readers import readColmapSceneInfo
    sig = inspect.signature(readColmapSceneInfo)
    print(f"Full signature: {sig}")
    print()
    print("按 render_single.py 的位置策略喂参 (source_path='<SRC>', img_dir='<IMG>'):")
    for i, (k, p) in enumerate(sig.parameters.items()):
        default = p.default if p.default != inspect.Parameter.empty else "<REQUIRED>"
        if i == 0:                    fed = "<SRC>"
        elif i == 1:                  fed = "<IMG>"
        elif k == "eval":             fed = "True"
        elif k == "train_test_exp":   fed = "False"
        else:                         fed = f"default({default})"
        warn = ""
        if k in ("eval", "train_test_exp") and i < 2:
            warn = "  <<< WARNING: positional slot collides with keyword logic"
        print(f"  [{i}] {k:<20} default={default!s:<20} -> fed: {fed}{warn}")


# ---------- Probe 2 ----------
def probe_native_gt(scene):
    sec(f"PROBE 2 — Native GT manifest [{scene}]")
    cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}")
    candidates = [
        os.path.join(cell, "gt_test_30000"),
        os.path.join(cell, "renders_test_30000", "gt"),
        os.path.join(cell, "test", "ours_30000", "gt"),
    ]
    gt_dir = next((c for c in candidates if os.path.isdir(c)), None)
    if gt_dir is None:
        print(f"[!] GT dir not found. Tried: {candidates}")
        return None
    files = sorted([f for f in os.listdir(gt_dir) if f.lower().endswith(('.jpg', '.png', '.jpeg'))])
    print(f"GT dir : {gt_dir}")
    print(f"Count  : {len(files)}")
    print(f"First 5: {files[:5]}")
    print(f"Last  3: {files[-3:]}")
    if files:
        img = Image.open(os.path.join(gt_dir, files[0]))
        print(f"Sample dims (first GT): {img.size}  (W x H)")
    return set(os.path.splitext(f)[0] for f in files)


# ---------- Probe 3 ----------
def probe_our_split(scene):
    sec(f"PROBE 3 — Our eval=True split [{scene}]")
    from scene.dataset_readers import readColmapSceneInfo
    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"
    print(f"Picked img_dir: {img_dir}")
    sig, args_list = build_args(source_path, img_dir)
    print(f"Args positionally passed:")
    for i, (k, _) in enumerate(sig.parameters.items()):
        print(f"  [{i}] {k} = {args_list[i]!r}")
    try:
        scene_info = readColmapSceneInfo(*args_list)
    except Exception as e:
        print(f"[!] readColmapSceneInfo crashed: {type(e).__name__}: {e}")
        return None
    tr, te = scene_info.train_cameras, scene_info.test_cameras
    print(f"Train cams : {len(tr)}")
    print(f"Test  cams : {len(te)}")
    if te:
        names = sorted(c.image_name for c in te)
        print(f"First 5 test names: {names[:5]}")
        print(f"Last  3 test names: {names[-3:]}")
        c0 = te[0]
        print(f"First test cam dims: W={c0.width} H={c0.height}  FovX={c0.FovX:.4f} FovY={c0.FovY:.4f}")
    return set(os.path.splitext(c.image_name)[0] for c in te)


# ---------- Probe 4 ----------
def probe_split_diff(scene, native_set, ours_set):
    sec(f"PROBE 4 — Split diff [{scene}]")
    if native_set is None or ours_set is None:
        print("[!] skipped (missing input)")
        return
    only_n = sorted(native_set - ours_set)
    only_o = sorted(ours_set - native_set)
    print(f"|Native|={len(native_set)}  |Ours|={len(ours_set)}  |intersection|={len(native_set & ours_set)}")
    print(f"In Native only ({len(only_n)}): {only_n[:5]}")
    print(f"In Ours   only ({len(only_o)}): {only_o[:5]}")
    if not only_n and not only_o:
        print("[OK] Test split fully aligned.")
    else:
        print("[!!] SPLIT MISMATCH — eval=True 没把我们带到和 Native 一样的测试集上")


# ---------- Probe 5 ----------
def probe_camera_model(scene):
    sec(f"PROBE 5 — COLMAP camera model [{scene}]")
    src = os.path.join(DATASET_ROOT, scene, "sparse", "0")
    txt, binf = os.path.join(src, "cameras.txt"), os.path.join(src, "cameras.bin")
    if os.path.exists(txt):
        with open(txt) as f:
            for line in f:
                if line.startswith('#') or not line.strip(): continue
                p = line.strip().split()
                print(f"  id={p[0]} model={p[1]} size={p[2]}x{p[3]} params={p[4:]}")
                break
    elif os.path.exists(binf):
        MODELS = {0:"SIMPLE_PINHOLE",1:"PINHOLE",2:"SIMPLE_RADIAL",3:"RADIAL",
                  4:"OPENCV",5:"OPENCV_FISHEYE",6:"FULL_OPENCV",7:"FOV",
                  8:"SIMPLE_RADIAL_FISHEYE",9:"RADIAL_FISHEYE",10:"THIN_PRISM_FISHEYE"}
        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 = MODELS.get(model_id, f"UNKNOWN({model_id})")
            print(f"  cameras.bin count={num}  id={cam_id}  model={model_name}  size={width}x{height}")
            if model_name not in ("SIMPLE_PINHOLE", "PINHOLE"):
                print(f"  [!!] Non-pinhole model — 畸变系数没被 gsplat 内参 K 建模！")
    else:
        print(f"[!] no cameras.txt or cameras.bin under {src}")


# ---------- Probe 6 ----------
def probe_exif(scene):
    sec(f"PROBE 6 — EXIF orientation [{scene}]")
    for folder in ["images", "images_2"]:
        d = os.path.join(DATASET_ROOT, scene, folder)
        if not os.path.isdir(d):
            continue
        files = sorted([f for f in os.listdir(d) if f.lower().endswith(('.jpg', '.png', '.jpeg'))])
        if not files:
            continue
        img = Image.open(os.path.join(d, files[0]))
        raw = img.size
        tag = None
        try:
            ex = img._getexif()
            if ex: tag = ex.get(274, None)   # 274 = Orientation
        except Exception:
            pass
        rot = ImageOps.exif_transpose(img).size
        flag = "[FLIPPED by exif_transpose]" if raw != rot else ""
        print(f"  {folder:<10} {files[0]:<30} raw={raw} exif_tag={tag} after={rot} {flag}")


# ---------- Probe 7 ----------
def probe_ply_z(scene):
    sec(f"PROBE 7 — PLY Z distribution in camera space [{scene}]")
    ply = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}",
                       "point_cloud", "iteration_30000", "point_cloud.ply")
    if not os.path.exists(ply):
        print(f"[!] PLY missing: {ply}")
        return
    v = PlyData.read(ply)['vertex']
    means = np.stack((v['x'], v['y'], v['z']), axis=-1).astype(np.float32)
    print(f"PLY gaussians : {len(means)}")
    print(f"World mean    : {means.mean(0)}")
    print(f"World std     : {means.std(0)}")

    from scene.dataset_readers import readColmapSceneInfo
    from utils.graphics_utils import getWorld2View2
    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"
    _, args_list = build_args(source_path, img_dir)
    te = readColmapSceneInfo(*args_list).test_cameras
    if not te:
        print("[!] no test cams")
        return
    c = te[0]
    W2V = getWorld2View2(np.array(c.R), np.array(c.T)).astype(np.float32)
    xyz_h = np.concatenate([means, np.ones((len(means), 1), dtype=np.float32)], axis=1)
    z = (xyz_h @ W2V.T)[:, 2]
    print(f"Viewed from test[0] = {c.image_name}")
    print(f"Z percentiles (camera +Z forward):")
    for q in [1, 5, 50, 95, 99, 99.9, 100]:
        print(f"  {q:5.1f}% : {np.percentile(z, q):10.2f}")
    print(f"Gaussians Z>50  : {(z>50).sum():>8} ({100*(z>50).mean():.2f}%)")
    print(f"Gaussians Z>100 : {(z>100).sum():>8} ({100*(z>100).mean():.2f}%)")
    print(f"Gaussians Z>500 : {(z>500).sum():>8} ({100*(z>500).mean():.2f}%)")


def main():
    probe_signature()
    for scene, label in SCENES:
        print(f"\n\n{'#'*70}\n# SCENE: {scene}  [{label}]\n{'#'*70}")
        native = probe_native_gt(scene)
        ours   = probe_our_split(scene)
        probe_split_diff(scene, native, ours)
        probe_camera_model(scene)
        probe_exif(scene)
        probe_ply_z(scene)


if __name__ == "__main__":
    main()