#!/usr/bin/env python
"""Ranking consistency probe — DSLR-FAIL scenes × golden methods.
Decides whether the ~1.2 dB systematic offset still preserves within-benchmark
method ordering.
"""
import os, sys, inspect, math, json, time
import numpy as np
import torch
from PIL import Image, ImageOps
from plyfile import PlyData

sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official')
import gsplat
try:
    from scipy.stats import spearmanr
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False

DATASET_ROOT = "/root/autodl-tmp/dataset/tnt"
OUTPUT_ROOT  = "/root/autodl-tmp/SplatAtlas/outputs"

METHODS = ["vanilla_3dgs", "erankgs", "ges", "lightgaussian", "opti3dgs",
           "reactgs", "steepgs", "absgs", "gaussianpro", "minisplatting", "pixelgs"]
SCENES  = ["palace", "lighthouse", "francis", "temple", "auditorium"]


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


def find_cell(method, scene):
    target = f"{method}_{scene}".lower()
    if not os.path.isdir(OUTPUT_ROOT):
        return None
    for d in os.listdir(OUTPUT_ROOT):
        if d.lower() == target:
            return os.path.join(OUTPUT_ROOT, d)
    return None


def load_ply_tensors(ply_path, device):
    v = PlyData.read(ply_path)['vertex']
    props = v.data.dtype.names
    def t32(x): return torch.tensor(x, dtype=torch.float32, device=device)
    means     = t32(np.stack((v['x'], v['y'], v['z']), -1))
    quats     = t32(np.stack((v['rot_0'], v['rot_1'], v['rot_2'], v['rot_3']), -1))
    scales    = torch.exp(t32(np.stack((v['scale_0'], v['scale_1'], v['scale_2']), -1)))
    opacities = torch.sigmoid(t32(np.asarray(v['opacity'])))
    n_rest = sum(1 for p in props if p.startswith('f_rest_'))
    sh_per_ch = n_rest // 3
    sh_degree_map = {0: 0, 3: 1, 8: 2, 15: 3}
    sh_degree = sh_degree_map.get(sh_per_ch, min(3, int(math.sqrt(sh_per_ch + 1)) - 1))
    f_dc = t32(np.stack((v['f_dc_0'], v['f_dc_1'], v['f_dc_2']), -1)).unsqueeze(1)
    if n_rest > 0:
        f_rest = t32(np.stack([v[f'f_rest_{i}'] for i in range(n_rest)], -1))
        f_rest = f_rest.view(-1, 3, sh_per_ch).transpose(1, 2)
        shs = torch.cat([f_dc, f_rest], dim=1)
    else:
        shs = f_dc
    return means, quats, scales, opacities, shs, sh_degree


def render_cell(method, scene, scene_info, device):
    from utils.graphics_utils import getWorld2View2
    from utils.general_utils import PILtoTorch

    cell = find_cell(method, scene)
    if cell is None:
        return None, None, "NO_CELL"

    bp = os.path.join(cell, "metrics_test_iter30000.json")
    native = None
    if os.path.exists(bp):
        bd = json.load(open(bp))
        native = bd.get("photometric", {}).get("PSNR")
        if native is None:
            native = bd.get("PSNR")

    ply_path = os.path.join(cell, "point_cloud", "iteration_30000", "point_cloud.ply")
    if not os.path.exists(ply_path):
        return None, native, "NO_PLY"

    try:
        means, quats, scales, opacities, shs, sh_degree = load_ply_tensors(ply_path, device)
    except Exception as e:
        return None, native, f"PLY_ERR:{type(e).__name__}"

    source_path = os.path.join(DATASET_ROOT, scene)
    img_dir = "images_2"
    resolution = 2
    test_cams = scene_info.test_cameras
    bg = torch.tensor([0., 0., 0.], device=device)

    psnrs = []
    try:
        for c in test_cams:
            w = int(round(c.width  / resolution))
            h = int(round(c.height / resolution))
            viewmat = torch.tensor(getWorld2View2(np.array(c.R), np.array(c.T)),
                                   dtype=torch.float32, device=device).unsqueeze(0)
            fx = w / (2 * math.tan(c.FovX / 2))
            fy = h / (2 * math.tan(c.FovY / 2))
            K = torch.tensor([[fx, 0, w/2], [0, fy, h/2], [0, 0, 1]],
                             dtype=torch.float32, device=device).unsqueeze(0)
            with torch.no_grad():
                colors, _, _ = gsplat.rasterization(
                    means=means, quats=quats, scales=scales, opacities=opacities,
                    colors=shs, viewmats=viewmat, Ks=K, width=w, height=h,
                    sh_degree=sh_degree, packed=True, render_mode='RGB',
                    backgrounds=bg.unsqueeze(0),
                    rasterize_mode='classic')
            render = colors[0].clamp(0, 1)
            pil = Image.open(os.path.join(source_path, img_dir, c.image_name))
            pil = ImageOps.exif_transpose(pil)
            gt_chw = PILtoTorch(pil, (w, h)).to(device)
            if gt_chw.shape[0] == 4: gt_chw = gt_chw[:3]
            gt = gt_chw.permute(1, 2, 0)
            mse = ((render - gt) ** 2).mean().item()
            psnrs.append(10 * math.log10(1.0 / max(mse, 1e-10)))
    except Exception as e:
        return None, native, f"RENDER_ERR:{type(e).__name__}"

    del means, quats, scales, opacities, shs
    torch.cuda.empty_cache()
    return float(np.mean(psnrs)), native, "OK"


def print_matrix(mat, methods, scenes, title, fmt="{:>10.3f}"):
    print(f"\n  {title}")
    header = "  " + " " * 16 + " ".join(f"{s[:10]:>10}" for s in scenes)
    print(header)
    for mi, method in enumerate(methods):
        row = " ".join(
            (fmt.format(mat[mi, si]) if not np.isnan(mat[mi, si]) else f"{'—':>10}")
            for si in range(len(scenes)))
        print(f"  {method:>16}  {row}")


def rank_from_psnr(arr, valid_mask):
    """Higher PSNR → rank 1. NaN / invalid → NaN."""
    ranks = np.full_like(arr, np.nan, dtype=float)
    valid_idx = np.where(valid_mask)[0]
    if len(valid_idx) == 0:
        return ranks
    vals = arr[valid_idx]
    order = np.argsort(-vals)
    for r, o in enumerate(order):
        ranks[valid_idx[o]] = r + 1
    return ranks


def spearman_rho(a, b):
    """Naive Spearman; fallback if scipy not available."""
    if HAS_SCIPY:
        rho, p = spearmanr(a, b)
        return rho, p
    ra = np.argsort(np.argsort(a))
    rb = np.argsort(np.argsort(b))
    mean_a, mean_b = ra.mean(), rb.mean()
    num = ((ra - mean_a) * (rb - mean_b)).sum()
    den = math.sqrt(((ra - mean_a) ** 2).sum() * ((rb - mean_b) ** 2).sum())
    return (num / den if den > 0 else float('nan'), float('nan'))


def main():
    from scene.dataset_readers import readColmapSceneInfo
    device = torch.device("cuda")

    print("Preloading scene infos...")
    scene_infos = {}
    for sc in SCENES:
        source_path = os.path.join(DATASET_ROOT, sc)
        if not os.path.exists(source_path):
            print(f"  [!] dataset path missing: {source_path}")
            continue
        scene_infos[sc] = readColmapSceneInfo(*build_args(source_path, "images_2"))
        print(f"  [preloaded] {sc}: {len(scene_infos[sc].test_cameras)} test cams")

    ours_mat   = np.full((len(METHODS), len(SCENES)), np.nan)
    native_mat = np.full((len(METHODS), len(SCENES)), np.nan)
    status_mat = [["-"] * len(SCENES) for _ in METHODS]

    print("\nRunning ...\n")
    t_all = time.time()
    for mi, method in enumerate(METHODS):
        for si, scene in enumerate(SCENES):
            if scene not in scene_infos:
                status_mat[mi][si] = "NOSRC"
                continue
            t0 = time.time()
            ours, native, status = render_cell(method, scene, scene_infos[scene], device)
            dt = time.time() - t0
            if ours is not None:   ours_mat[mi, si]   = ours
            if native is not None: native_mat[mi, si] = native
            status_mat[mi][si] = status
            delta = (ours - native) if (ours is not None and native is not None) else float('nan')
            ours_s   = f"{ours:7.3f}"   if ours   is not None else "    ×  "
            native_s = f"{native:7.3f}" if native is not None else "    ×  "
            delta_s  = f"{delta:+7.3f}" if not (isinstance(delta, float) and math.isnan(delta)) else "   ×   "
            print(f"  [{method:<14} {scene:<11}] ours={ours_s}  native={native_s}  Δ={delta_s}  "
                  f"({dt:5.1f}s  {status})")
    print(f"\nTotal time: {time.time() - t_all:.1f}s")

    # === Report ===
    print("\n" + "=" * 90)
    print("  RANKING CONSISTENCY REPORT")
    print("=" * 90)

    print_matrix(ours_mat,   METHODS, SCENES, "OURS PSNR (gsplat classic)",  "{:>10.3f}")
    print_matrix(native_mat, METHODS, SCENES, "NATIVE baseline PSNR",        "{:>10.3f}")
    delta_mat = ours_mat - native_mat
    print_matrix(delta_mat,  METHODS, SCENES, "Δ (ours - native)",           "{:>+10.3f}")

    valid_delta = delta_mat[~np.isnan(delta_mat)]
    valid_ours  = ours_mat[~np.isnan(ours_mat)]
    print(f"\n  === Δ stats (N={len(valid_delta)}) ===")
    print(f"    mean    = {valid_delta.mean():+.4f}   std   = {valid_delta.std():.4f}")
    print(f"    min     = {valid_delta.min():+.4f}   max   = {valid_delta.max():+.4f}")
    print(f"    |Δ|>1dB : {(np.abs(valid_delta) > 1).sum()}/{len(valid_delta)}")
    print(f"    |Δ|>2dB : {(np.abs(valid_delta) > 2).sum()}/{len(valid_delta)}")
    print(f"    |Δ|>3dB : {(np.abs(valid_delta) > 3).sum()}/{len(valid_delta)}")

    print(f"\n  === OURS absolute value distribution ===")
    print(f"    min     = {valid_ours.min():.3f}   max   = {valid_ours.max():.3f}")
    print(f"    <10 dB  : {(valid_ours < 10).sum()}/{len(valid_ours)}  [CATASTROPHIC]")
    print(f"    <15 dB  : {(valid_ours < 15).sum()}/{len(valid_ours)}")
    print(f"    <20 dB  : {(valid_ours < 20).sum()}/{len(valid_ours)}")

    # Flag catastrophic cells
    cat_cells = []
    for mi in range(len(METHODS)):
        for si in range(len(SCENES)):
            if not np.isnan(ours_mat[mi, si]) and ours_mat[mi, si] < 15:
                cat_cells.append((METHODS[mi], SCENES[si], ours_mat[mi, si],
                                  native_mat[mi, si], delta_mat[mi, si]))
    if cat_cells:
        print(f"\n  === Cells with OURS < 15 dB ===")
        for m, s, o, n, d in sorted(cat_cells, key=lambda x: x[2]):
            print(f"    {m:>16} × {s:<12}  ours={o:7.3f}  native={n:7.3f}  Δ={d:+.3f}")

    # Per-scene ranking consistency
    print(f"\n  === Per-scene ranking (Spearman ρ, higher PSNR = rank 1) ===")
    print(f"  {'scene':>12}  {'ρ':>8}  {'p':>8}  {'N':>3}   concordant_pairs")
    for si, scene in enumerate(SCENES):
        valid = ~np.isnan(ours_mat[:, si]) & ~np.isnan(native_mat[:, si])
        n = valid.sum()
        if n < 3:
            print(f"  {scene:>12}  {'—':>8}  {'—':>8}  {n:>3}")
            continue
        rho, p = spearman_rho(ours_mat[valid, si], native_mat[valid, si])
        r_ours   = rank_from_psnr(ours_mat[:, si],   valid)
        r_native = rank_from_psnr(native_mat[:, si], valid)
        # 相邻排名差异
        rank_diff = np.abs(r_ours - r_native)
        max_rd = np.nanmax(rank_diff)
        p_str = f"{p:.4f}" if not math.isnan(p) else "n/a"
        print(f"  {scene:>12}  {rho:>+8.4f}  {p_str:>8}  {n:>3}   max_rank_shift={max_rd:.0f}")

    # Per-method consistency (same method across scenes)
    print(f"\n  === Per-method ranking across {len(SCENES)} scenes ===")
    for mi, method in enumerate(METHODS):
        valid = ~np.isnan(ours_mat[mi, :]) & ~np.isnan(native_mat[mi, :])
        n = valid.sum()
        if n < 3: continue
        rho, _ = spearman_rho(ours_mat[mi, valid], native_mat[mi, valid])
        print(f"    {method:>16}  ρ={rho:+.4f}  (N={n})")

    # Save JSON
    out_json = "/root/autodl-tmp/SplatAtlas/scripts/phase1_validation/ranking_consistency.json"
    out = {
        "methods": METHODS, "scenes": SCENES,
        "ours_psnr":   ours_mat.tolist(),
        "native_psnr": native_mat.tolist(),
        "delta":       delta_mat.tolist(),
        "status":      status_mat,
    }
    with open(out_json, "w") as f:
        json.dump(out, f, indent=2, default=str)
    print(f"\n  Saved: {out_json}")

    # Final verdict
    print(f"\n  === VERDICT ===")
    ok_abs      = (valid_ours >= 15).all()
    ok_delta    = np.abs(valid_delta).mean() < 2.0
    n_rho_high  = 0
    if HAS_SCIPY or not HAS_SCIPY:
        for si, scene in enumerate(SCENES):
            valid = ~np.isnan(ours_mat[:, si]) & ~np.isnan(native_mat[:, si])
            if valid.sum() >= 3:
                rho, _ = spearman_rho(ours_mat[valid, si], native_mat[valid, si])
                if rho > 0.8:
                    n_rho_high += 1
    print(f"    OURS absolute ≥15dB:        {'YES' if ok_abs else 'NO'}")
    print(f"    |Δ|mean < 2dB:              {'YES' if ok_delta else 'NO'}")
    print(f"    Scenes with ρ>0.8:          {n_rho_high}/{len(SCENES)}")
    if ok_abs and n_rho_high >= len(SCENES) - 1:
        print(f"    → Benchmark 内部一致性良好。可以收工。")
    else:
        print(f"    → 有异常，需要看具体哪些 cell 不稳。")


if __name__ == "__main__":
    main()