import os
import sys
import subprocess
import json
import csv
import glob
from tqdm import tqdm

BASE_DIR = "/root/autodl-tmp/SplatAtlas"
OUT_BASE = os.path.join(BASE_DIR, "outputs")
CSV_PATH = os.path.join(BASE_DIR, "scripts/phase1_validation/validation_results_anchor.csv")
ERR_LOG  = os.path.join(BASE_DIR, "scripts/phase1_validation/batch_errors.log")

METHODS = [
    "3dgsmcmc", "analyticsplatting", "atomgs", "coadaptgs", "conegs", 
    "erankgs", "ges", "ghap", "gof", "gslpm", "lightgaussian", "opti3dgs", 
    "pgsr", "reactgs", "steepgs", "vanilla_3dgs", "absgs", "gaussianpro", 
    "minisplatting", "pixelgs"
]

SCENES = {
    "bonsai":    {"dataset": "mip360_indoor",  "res": 2, "bg": "0,0,0", "source": "/root/autodl-tmp/dataset/360/bonsai"},
    "garden":    {"dataset": "mip360_outdoor", "res": 4, "bg": "0,0,0", "source": "/root/autodl-tmp/dataset/360/garden"},
    "truck":     {"dataset": "tnt",            "res": 1, "bg": "0,0,0", "source": "/root/autodl-tmp/dataset/tnt/truck"},
    "drjohnson": {"dataset": "db",             "res": 1, "bg": "0,0,0", "source": "/root/autodl-tmp/dataset/deepblending_clean/DrJohnson"}
}

def log_error(cell_name, phase, stdout, stderr):
    with open(ERR_LOG, "a") as f:
        f.write(f"\n{'='*50}\n[{phase} ERROR] Cell: {cell_name}\n")
        if stdout: f.write(f"--- STDOUT (last 20 lines) ---\n" + "\n".join(stdout.strip().split("\n")[-20:]) + "\n")
        if stderr: f.write(f"--- STDERR (last 20 lines) ---\n" + "\n".join(stderr.strip().split("\n")[-20:]) + "\n")
        f.write(f"{'='*50}\n")

def get_actual_cell_dir(method, scene):
    target = f"{method}_{scene}".lower()
    if os.path.exists(OUT_BASE):
        for d in os.listdir(OUT_BASE):
            if d.lower() == target:
                return os.path.join(OUT_BASE, d)
    return None

def phase_a_native(method, scene_key, scene_cfg, cell_dir):
    json_path = os.path.join(cell_dir, "metrics_test_iter30000.json")
    if os.path.exists(json_path):
        try:
            with open(json_path, 'r') as f:
                data = json.load(f)
                if "PSNR" in data or ("photometric" in data and "PSNR" in data["photometric"]):
                    return True
        except:
            pass

    # Step 1: Render
    cmd_render = [
        "python", "scripts/main_render.py", "--method", method,
        "--source_path", scene_cfg["source"], "--model_path", cell_dir,
        "--iteration", "30000", "--resolution", str(scene_cfg["res"])
    ]
    res = subprocess.run(cmd_render, cwd=BASE_DIR, capture_output=True, text=True)
    if res.returncode != 0:
        log_error(f"{method}_{scene_key}", "PHASE_A_RENDER", res.stdout, res.stderr)
        return False

    # Step 2: Compute Metrics
    py_metric = f"""
import sys, json, os
sys.path.append("ufd_evalkit")
from photometric import compute_photometric_metrics
r_dir = "{cell_dir}/renders_test_30000/renders"
g_dir = "{cell_dir}/renders_test_30000/gt"
if not os.path.exists(r_dir):
    r_dir = "{cell_dir}/renders_test_30000"
    g_dir = "{cell_dir}/gt_test_30000"
metrics = compute_photometric_metrics(r_dir, g_dir)
with open("{json_path}", "w") as f:
    json.dump({{"PSNR": metrics.get("PSNR", 0)}}, f, indent=4)
"""
    res2 = subprocess.run(["python", "-c", py_metric], cwd=BASE_DIR, capture_output=True, text=True)
    if res2.returncode != 0:
        log_error(f"{method}_{scene_key}", "PHASE_A_METRIC", res2.stdout, res2.stderr)
        return False
    return True

def phase_b_gsplat(ply_path, scene_cfg, cell_dir, cell_name):
    out_dir = os.path.join(OUT_BASE, "phase1_validation", cell_name)
    os.makedirs(out_dir, exist_ok=True)
    cmd = [
        "python", "scripts/phase1_validation/render_single.py",
        "--ply_path", ply_path, "--source_path", scene_cfg["source"],
        "--model_path", cell_dir, "--output_dir", out_dir,
        "--resolution", str(scene_cfg["res"]), "--bg_color", scene_cfg["bg"]
    ]
    res = subprocess.run(cmd, cwd=BASE_DIR, capture_output=True, text=True)
    
    if res.returncode != 0:
        notes = "ERROR_B"
        if "CUDA out of memory" in res.stderr or "CUDA out of memory" in res.stdout:
            notes = "OOM"
        elif "plyfile" in res.stderr and "format" in res.stderr:
            notes = "PLY format incompatible"
        log_error(cell_name, "PHASE_B_GSPLAT", res.stdout, res.stderr)
        return None, None, None, notes

    native_psnr, gsplat_psnr, delta = None, None, None
    for line in res.stdout.split('\n'):
        if "Mean PSNR (ours/gsplat)" in line:
            try: gsplat_psnr = float(line.split(":")[1].strip().split()[0])
            except: pass
        if "Mean PSNR (native baseline)" in line:
            try: native_psnr = float(line.split(":")[1].strip().split()[0])
            except: pass
        if "Delta" in line and "dB" in line and "STATUS" not in line:
            try: delta = float(line.split(":")[1].strip().split()[0])
            except: pass
            
    # Fallback to json if regex fails
    if gsplat_psnr is None:
        try:
            with open(os.path.join(out_dir, "psnr_results.json"), "r") as f:
                d = json.load(f)
                gsplat_psnr = float(np.mean(list(d.values())))
        except: pass

    notes = ""
    if gsplat_psnr is not None and gsplat_psnr < 15.0:
        notes = "quantized PLY? low PSNR"

    return native_psnr, gsplat_psnr, delta, notes

def main():
    print("=== SplatAtlas Phase 1 Task B: Anchor Batch Validation ===")
    
    completed = set()
    write_header = not os.path.exists(CSV_PATH)
    if not write_header:
        with open(CSV_PATH, "r") as f:
            reader = csv.DictReader(f)
            for row in reader:
                completed.add((row['method'], row['scene']))

    total_cells = len(METHODS) * len(SCENES)
    pbar = tqdm(total=total_cells, desc="Validating Cells", unit="cell")

    with open(CSV_PATH, "a", newline='') as f:
        writer = csv.writer(f)
        if write_header:
            writer.writerow(["method", "scene", "dataset", "resolution", "has_ply", "native_psnr", "gsplat_psnr", "delta_db", "status", "notes"])

        for method in METHODS:
            for scene, cfg in SCENES.items():
                if (method, scene) in completed:
                    pbar.update(1)
                    continue

                cell_dir = get_actual_cell_dir(method, scene)
                if not cell_dir:
                    writer.writerow([method, scene, cfg["dataset"], cfg["res"], "False", "", "", "", "SKIP", "Directory missing"])
                    f.flush()
                    pbar.set_postfix_str(f"[{method}/{scene}] SKIP (No Dir)")
                    pbar.update(1)
                    continue

                ply_path = os.path.join(cell_dir, "point_cloud/iteration_30000/point_cloud.ply")
                if not os.path.exists(ply_path):
                    writer.writerow([method, scene, cfg["dataset"], cfg["res"], "False", "", "", "", "SKIP", "PLY missing"])
                    f.flush()
                    pbar.set_postfix_str(f"[{method}/{scene}] SKIP (No PLY)")
                    pbar.update(1)
                    continue

                cell_name = os.path.basename(cell_dir)
                
                # Phase A
                if not phase_a_native(method, scene, cfg, cell_dir):
                    writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", "", "", "", "ERROR_A", "Native render/eval failed"])
                    f.flush()
                    pbar.set_postfix_str(f"[{method}/{scene}] ERROR_A")
                    pbar.update(1)
                    continue

                # Phase B
                native, gsplat, delta, notes = phase_b_gsplat(ply_path, cfg, cell_dir, cell_name)
                
                if native is None and gsplat is None:
                    writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", "", "", "", "ERROR_B", notes])
                    f.flush()
                    pbar.set_postfix_str(f"[{method}/{scene}] ERROR_B")
                else:
                    if delta is None and native is not None and gsplat is not None:
                        delta = gsplat - native
                    
                    status = "FAIL"
                    if delta is not None:
                        if abs(delta) < 0.3: status = "PASS"
                        elif abs(delta) < 0.5: status = "WARN"
                        
                    writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", 
                                     f"{native:.4f}" if native else "", 
                                     f"{gsplat:.4f}" if gsplat else "", 
                                     f"{delta:.4f}" if delta else "", 
                                     status, notes])
                    f.flush()
                    pbar.set_postfix_str(f"[{method}/{scene}] {status} Δ={delta:+.2f}dB")
                
                pbar.update(1)
    
    pbar.close()

    # Generate Summary
    stats = {"PASS": 0, "WARN": 0, "FAIL": 0, "SKIP": 0, "ERROR": 0}
    method_summary = {m: {"PASS": 0, "total_run": 0, "fails": []} for m in METHODS}
    
    with open(CSV_PATH, "r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            st = row['status']
            if st.startswith("ERROR"): stats["ERROR"] += 1
            elif st in stats: stats[st] += 1
            
            m = row['method']
            if m in method_summary and st not in ("SKIP", "ERROR_A"):
                method_summary[m]["total_run"] += 1
                if st == "PASS": method_summary[m]["PASS"] += 1
                elif st in ("FAIL", "WARN", "ERROR_B"):
                    method_summary[m]["fails"].append(f"{row['scene']} Δ={row['delta_db']} ({st})")

    total = sum(stats.values())
    print("\n" + "="*40)
    print("==== BATCH SUMMARY ====")
    print(f"  PASS:  {stats['PASS']}/{total}")
    print(f"  WARN:  {stats['WARN']}")
    print(f"  FAIL:  {stats['FAIL']}")
    print(f"  SKIP:  {stats['SKIP']}")
    print(f"  ERROR: {stats['ERROR']}")
    print("\nFAIL/WARN Details by Method:")
    
    all_pass = []
    all_fail = []
    mixed = []
    
    for m, d in method_summary.items():
        if d["total_run"] == 0: continue
        if d["PASS"] == d["total_run"]:
            all_pass.append(m)
        elif d["PASS"] == 0:
            all_fail.append(m)
            print(f"  {m}: ALL FAIL/ERROR -> " + ", ".join(d["fails"]))
        else:
            mixed.append(m)
            print(f"  {m}: Mixed -> " + ", ".join(d["fails"]))

    print("\nMethod Classification:")
    print(f"  Fully PASS ({len(all_pass)}): " + ", ".join(all_pass))
    print(f"  Fully FAIL/ERR ({len(all_fail)}) [Needs DataLoader Branch]: " + ", ".join(all_fail))
    print(f"  Mixed ({len(mixed)}) [Scene-dependent issues]: " + ", ".join(mixed))
    print("========================================")
    print("BATCH ANCHOR DONE")
    print(f"See: {CSV_PATH}")

if __name__ == "__main__":
    main()