import os, json, numpy as np

OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs/error_responsibility"
PROF_DIR = os.path.join(OUTPUT_ROOT, "_profiles")

METHODS = ['vanilla_3dgs', 'analyticsplatting', 'erankgs', 'ges', 'lightgaussian', 'minisplatting', 'opti3dgs', 'pgsr', 'steepgs']
SCENES = ['bicycle','bonsai','counter','flowers','garden','kitchen','room','stump','treehill',
          'auditorium','ballroom','barn','caterpillar','courtroom','lighthouse','museum','palace','playground','temple','train','truck',
          'DrJohnson','Playroom',
          'Chair','Drums','Ficus','Hotdog','Lego','Materials','Mic','Ship']
METRICS = ["psnr_full", "sh_net_effect", "sh_corruption_rate", "opacity_net_effect", "coverage", "coverage_error_fraction", "residual_error"]

data = {m: {} for m in METHODS}
status_counts = {m: {"PASS": 0, "WARN": 0, "FAIL": 0} for m in METHODS}
failed_cells = []

# --- 1. Load Data & Status Statistics ---
for m in METHODS:
    for s in SCENES:
        p = os.path.join(OUTPUT_ROOT, f"{m}_{s}", "counterfactual_summary.json")
        if os.path.exists(p):
            with open(p) as f:
                d = json.load(f)
                data[m][s] = d
                st = d.get("sanity_status", "FAIL")
                status_counts[m][st] = status_counts[m].get(st, 0) + 1
                if st == "FAIL":
                    failed_cells.append(f"{m}_{s}: {d.get('sanity_failures', [])}")

medians = {m: {} for m in METHODS}
for m in METHODS:
    for k in METRICS:
        vals = [data[m][s][k] for s in SCENES if s in data[m] and data[m][s].get(k) is not None]
        medians[m][k] = float(np.median(vals)) if vals else 0.0

# --- 2. Saturated Subset Analysis ---
# 寻找 PSNR 最接近的 4 个方法构成的饱和簇
sorted_by_psnr = sorted(METHODS, key=lambda x: medians[x]["psnr_full"])
best_window = []
min_psnr_range = 9999

for i in range(len(sorted_by_psnr) - 3):
    window = sorted_by_psnr[i:i+4]
    rng = medians[window[-1]]["psnr_full"] - medians[window[0]]["psnr_full"]
    if rng < min_psnr_range:
        min_psnr_range = rng
        best_window = window

def calc_iqr_stats(vals_list):
    if not vals_list: return None, None, None
    med = float(np.median(vals_list))
    abs_iqr = float(np.percentile(vals_list, 75) - np.percentile(vals_list, 25))
    rel_iqr = (abs_iqr / abs(med) * 100) if med != 0 else None
    return med, abs_iqr, rel_iqr

sat_psnr_vals = [medians[m]["psnr_full"] for m in best_window]
_, _, sat_psnr_rel_iqr = calc_iqr_stats(sat_psnr_vals)

sat_results = {"subset_methods": best_window, "psnr_range_db": min_psnr_range, "psnr_rel_iqr": sat_psnr_rel_iqr, "metrics": {}}
for k in METRICS:
    if k == "psnr_full": continue
    vals = [medians[m][k] for m in best_window]
    _, _, rel_iqr = calc_iqr_stats(vals)
    sat_results["metrics"][k] = {
        "rel_iqr": rel_iqr, 
        "ratio_vs_psnr": (rel_iqr / sat_psnr_rel_iqr) if (rel_iqr is not None and sat_psnr_rel_iqr) else None
    }

with open(os.path.join(PROF_DIR, "psnr_saturated_subset_analysis.json"), "w") as f:
    json.dump(sat_results, f, indent=4)

# --- 3. Stratification (Delta Signs vs Vanilla) ---
strat_results = {}
for m in METHODS:
    if m == "vanilla_3dgs": continue
    strat_results[m] = {}
    for k in METRICS:
        if k == "psnr_full": continue
        delta = medians[m][k] - medians["vanilla_3dgs"][k]
        strat_results[m][k] = {"delta": delta, "sign": "+" if delta > 0 else ("-" if delta < 0 else "0")}

with open(os.path.join(PROF_DIR, "method_family_stratified.json"), "w") as f:
    json.dump(strat_results, f, indent=4)

# --- Output Report ---
print("=== 1. Execution Sanity Status ===")
print(f"{'Method':<20} | {'PASS':<6} | {'WARN':<6} | {'FAIL':<6}")
print("-" * 45)
for m in METHODS:
    print(f"{m:<20} | {status_counts[m]['PASS']:<6} | {status_counts[m]['WARN']:<6} | {status_counts[m]['FAIL']:<6}")
print(f"\nTotal FAILS across all {len(METHODS)*len(SCENES)} cells: {len(failed_cells)}")

print(f"\n=== 2. PSNR Saturated Subset Analysis ===")
print(f"Subset identified: {best_window}")
print(f"PSNR Relative IQR of this subset: {sat_psnr_rel_iqr:.2f}% (Range: {min_psnr_range:.2f} dB)")
print(f"\n{'Metric':<25} | {'Rel IQR (%)':<15} | {'Ratio vs PSNR':<15}")
print("-" * 60)
for k in METRICS:
    if k == "psnr_full": continue
    d = sat_results["metrics"][k]
    rel = f"{d['rel_iqr']:.2f}%" if d['rel_iqr'] else "N/A"
    rat = f"{d['ratio_vs_psnr']:.1f}x" if d['ratio_vs_psnr'] else "N/A"
    print(f"{k:<25} | {rel:<15} | {rat:<15}")

print(f"\n=== 3. Family Stratification (Sign vs Vanilla) ===")
print(f"{'Method':<20} | {'sh_net':<8} | {'sh_corr':<8} | {'opa_net':<8} | {'cov_err':<8} | {'resid':<8}")
print("-" * 70)
for m in METHODS:
    if m == "vanilla_3dgs": continue
    d = strat_results[m]
    print(f"{m:<20} | {d['sh_net_effect']['sign']:<8} | {d['sh_corruption_rate']['sign']:<8} | {d['opacity_net_effect']['sign']:<8} | {d['coverage_error_fraction']['sign']:<8} | {d['residual_error']['sign']:<8}")