#!/usr/bin/env python3
"""
Print Gaussian count, timing, and VRAM for all our completed experiments.

Data sources (in priority order):
  - Gaussian count  : PLY header  (element vertex N)
  - iter_time_ms    : TensorBoard  train/iter_time_ms  (last value)
  - peak_vram_GB    : TensorBoard  train/peak_vram_GB  (max across run)
  - timing fallback : timing.json cumulative_gpu_time_sec / iters

Usage:
    python scripts/print_efficiency.py
"""
import os
import json
import glob
import numpy as np

OUTPUTS_ROOT = "/root/autodl-tmp/SplatAtlas/outputs"
ITERS = 30000

KNOWN = ("vanillasgf", "vanilla_3dgs", "vanilla_logged", "vanilla_thresh",
         "absgs", "pixelgs", "edgeloss")

OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"}
INDOOR_360  = {"bonsai", "counter", "kitchen", "room"}
TNT = {
    "auditorium", "ballroom", "courtroom", "family", "francis", "horse",
    "lighthouse", "m60", "museum", "palace", "panther", "playground",
    "temple", "train"
}
ALL_SCENES = OUTDOOR_360 | INDOOR_360 | TNT


def is_ours(method):
    return any(method == p or method.startswith(p) for p in KNOWN)


def parse_dir(d):
    if d.startswith("vanilla_thresh_"):
        rest = d[len("vanilla_thresh_"):]
        for s in ALL_SCENES:
            if rest == s or rest.startswith(s + "_"):
                return "vanilla_thresh", s
        return None, None
    tok = d.split("_")
    for i in range(len(tok) - 1, 0, -1):
        sc = "_".join(tok[i:])
        me = "_".join(tok[:i])
        if sc in ALL_SCENES:
            return me, sc
    return None, None


def count_gaussians(ply_path):
    try:
        with open(ply_path, "rb") as f:
            for line in f:
                line = line.decode("ascii", errors="ignore").strip()
                if line.startswith("element vertex"):
                    return int(line.split()[-1])
                if line == "end_header":
                    break
    except Exception:
        pass
    return None


def read_tb(model_path):
    """Read last iter_time_ms and max peak_vram_GB from TensorBoard events."""
    try:
        from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
        ea = EventAccumulator(model_path, size_guidance={"scalars": 0})
        ea.Reload()
        tags = ea.Tags().get("scalars", [])

        it_ms = None
        for tag in ["train/iter_time_ms", "iter_time_ms"]:
            if tag in tags:
                vals = ea.Scalars(tag)
                if vals:
                    # mean of last 10% steps (stable phase)
                    n = max(1, len(vals) // 10)
                    it_ms = float(np.mean([v.value for v in vals[-n:]]))
                break

        vram = None
        for tag in ["train/peak_vram_GB", "peak_vram_GB"]:
            if tag in tags:
                vals = ea.Scalars(tag)
                if vals:
                    vram = float(max(v.value for v in vals))
                break

        return it_ms, vram
    except Exception:
        return None, None


def read_timing_json(model_path, iters):
    """Fallback: cumulative_gpu_time_sec from timing.json."""
    try:
        p = os.path.join(model_path, "timing.json")
        d = json.load(open(p))
        cgt = d.get("cumulative_gpu_time_sec")
        if cgt:
            return float(cgt) / iters * 1000   # ms/iter
    except Exception:
        pass
    return None


def main():
    plys = sorted(glob.glob(
        f"{OUTPUTS_ROOT}/*/point_cloud/iteration_{ITERS}/point_cloud.ply"))

    rows = []
    for ply_path in plys:
        parts = ply_path.split("/")
        dirname = parts[-4]
        method, scene = parse_dir(dirname)
        if method is None or not is_ours(method):
            continue

        model_path = os.path.join(OUTPUTS_ROOT, dirname)
        n_gauss = count_gaussians(ply_path)
        it_ms, vram = read_tb(model_path)
        if it_ms is None:
            it_ms = read_timing_json(model_path, ITERS)

        rows.append((dirname, method, scene, n_gauss, it_ms, vram))

    # ── Print ──────────────────────────────────────────────────────────────
    print(f"\n{'Dir':<42} {'#Gauss':>10} {'ms/iter':>9} {'PeakGB':>8}")
    print("-" * 73)

    by_method = {}
    for dirname, method, scene, n, t, v in rows:
        ns = f"{n:,}" if n else "-"
        ts = f"{t:.2f}" if t else "-"
        vs = f"{v:.3f}" if v else "-"
        print(f"{dirname:<42} {ns:>10} {ts:>9} {vs:>8}")
        by_method.setdefault(method, []).append((n, t, v))

    # ── Per-method averages ────────────────────────────────────────────────
    print(f"\n{'Method':<35} {'Avg #Gauss':>11} {'Avg ms/iter':>12} {'Avg PeakGB':>11} {'N':>4}")
    print("-" * 67)
    for m, vals in sorted(by_method.items()):
        ns = [x[0] for x in vals if x[0] is not None]
        ts = [x[1] for x in vals if x[1] is not None]
        vs = [x[2] for x in vals if x[2] is not None]
        nav = f"{int(np.mean(ns)):,}" if ns else "-"
        tav = f"{np.mean(ts):.2f}" if ts else "-"
        vav = f"{np.mean(vs):.3f}" if vs else "-"
        print(f"{m:<35} {nav:>11} {tav:>12} {vav:>11} {len(vals):>4}")


if __name__ == "__main__":
    main()