#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import csv
import math
from collections import defaultdict
from pathlib import Path

IN = Path("outputs/phase5a/task_vapre_splatatlas_repro.csv")
OUT = Path("outputs/phase1/per_cell_metrics_long.csv")
OUT.parent.mkdir(parents=True, exist_ok=True)

METHOD_COL = "method"
SCENE_COL = "scene"

METRICS = {
    "psnr_splat": "psnr",
    "ssim_splat": "ssim",
    "lpips_splat": "lpips",
}

EXPECTED_METHODS = [
    "3dgsmcmc", "absgs", "atomgs", "conegs", "ges", "ghap", "gslpm",
    "lapisgs", "reactgs", "vanilla_3dgs",
    "2dgs", "gaussian_surfel", "gof", "pgsr", "scaffoldgs",
    "coadaptgs", "erankgs", "gaussianpro", "minisplatting",
    "opti3dgs", "steepgs",
    "3dgs_dr", "analyticsplatting", "lod_gs", "mipsplatting",
    "pixelgs",
    "cdcgs", "hogs", "lightgaussian", "octree_gs", "trimgs",
]

EXPECTED_SCENES = [
    "auditorium", "ballroom", "barn", "caterpillar", "courtroom",
    "lighthouse", "museum", "palace", "playground", "temple", "train", "truck",
    "bicycle", "bonsai", "counter", "flowers", "garden", "kitchen", "room",
    "stump", "treehill",
    "chair", "drums", "ficus", "hotdog", "lego", "materials", "mic", "ship",
    "drjohnson", "playroom",
]

def parse_float(x):
    x = str(x).strip()
    if x == "" or x.lower() in {"nan", "none", "null", "na", "n/a", "---"}:
        return None
    try:
        v = float(x)
    except ValueError:
        return None
    if math.isnan(v) or math.isinf(v):
        return None
    return v

buckets = defaultdict(list)

with open(IN, "r", newline="", encoding="utf-8") as f:
    reader = csv.DictReader(f)
    fields = set(reader.fieldnames or [])

    required = {METHOD_COL, SCENE_COL, *METRICS.keys()}
    missing = required - fields
    if missing:
        raise ValueError(f"Missing columns in {IN}: {missing}. Actual columns={reader.fieldnames}")

    for row in reader:
        method = row[METHOD_COL].strip()
        scene = row[SCENE_COL].strip()

        for src_col, metric in METRICS.items():
            value = parse_float(row[src_col])
            if value is None:
                continue
            buckets[(method, scene, metric)].append(value)

rows = []
for (method, scene, metric), vals in sorted(buckets.items()):
    # If duplicate / multi-seed rows exist, average them.
    value = sum(vals) / len(vals)
    rows.append({
        "method_key": method,
        "scene": scene,
        "metric": metric,
        "value": f"{value:.10g}",
    })

with open(OUT, "w", newline="", encoding="utf-8") as f:
    writer = csv.DictWriter(f, fieldnames=["method_key", "scene", "metric", "value"])
    writer.writeheader()
    writer.writerows(rows)

methods = sorted({r["method_key"] for r in rows})
scenes = sorted({r["scene"] for r in rows})

print("Wrote:", OUT)
print("rows:", len(rows))
print("methods:", len(methods), methods)
print("scenes:", len(scenes), scenes)

for metric in ["psnr", "ssim", "lpips"]:
    cells = {(r["method_key"], r["scene"]) for r in rows if r["metric"] == metric}
    print(metric, "cells:", len(cells))

missing_methods = sorted(set(EXPECTED_METHODS) - set(methods))
extra_methods = sorted(set(methods) - set(EXPECTED_METHODS))
missing_scenes = sorted(set(EXPECTED_SCENES) - set(scenes))
extra_scenes = sorted(set(scenes) - set(EXPECTED_SCENES))

print("missing_methods:", missing_methods)
print("extra_methods:", extra_methods)
print("missing_scenes:", missing_scenes)
print("extra_scenes:", extra_scenes)

print("erankgs × lego:")
for r in rows:
    if r["method_key"] == "erankgs" and r["scene"].lower() == "lego":
        print(r)

if len(methods) != 31:
    raise AssertionError(f"Expected 31 methods, got {len(methods)}")
if len(scenes) != 31:
    raise AssertionError(f"Expected 31 scenes, got {len(scenes)}")
if missing_methods or extra_methods:
    raise AssertionError(f"Method mismatch: missing={missing_methods}, extra={extra_methods}")
if missing_scenes or extra_scenes:
    raise AssertionError(f"Scene mismatch: missing={missing_scenes}, extra={extra_scenes}")

expected_rows = 31 * 31 * 3
if len(rows) != expected_rows:
    print(f"[WARN] Expected {expected_rows} metric rows, got {len(rows)}.")
    print("[WARN] Missing metric values will be skipped in dataset means.")
    print("[WARN] If an entire method-dataset block is missing, it will appear as --- in Appendix E.")