"""Plot edge/flat densification-event ratio over training steps (Exp 2 curve)."""
import os, json, argparse
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--methods", nargs="+", default=["vanilla_logged","vanillasgf_logged"])
    ap.add_argument("--scenes", nargs="+", default=["francis","m60","panther","bonsai","bicycle","garden"])
    ap.add_argument("--outputs_root", default="/root/autodl-tmp/SplatAtlas/outputs")
    ap.add_argument("--out", default="/root/autodl-tmp/SplatAtlas/densification_timeseries.png")
    ap.add_argument("--bin_size", type=int, default=1000)
    args = ap.parse_args()

    fig, ax = plt.subplots(figsize=(8, 5))
    colors = {"vanilla_logged":"#888888", "vanillasgf_logged":"#1f77b4"}
    for method in args.methods:
        all_steps, all_ratios = [], []
        for scene in args.scenes:
            p = os.path.join(args.outputs_root, f"{method}_{scene}", "densification_events.json")
            if not os.path.exists(p): continue
            for e in json.load(open(p)):
                all_steps.append(e["step"])
                all_ratios.append(e["n_edge"] / max(e["n_flat"], 1))
        if not all_steps:
            print(f"no data for {method}"); continue
        steps, ratios = np.array(all_steps), np.array(all_ratios)
        bins = np.arange(0, int(steps.max())+args.bin_size, args.bin_size)
        bc, bm = [], []
        for i in range(len(bins)-1):
            m = (steps >= bins[i]) & (steps < bins[i+1])
            if m.sum() > 0:
                bc.append((bins[i]+bins[i+1])/2); bm.append(ratios[m].mean())
        ax.plot(bc, bm, label=method, marker='o', markersize=4, linewidth=1.5,
                color=colors.get(method))
    ax.axhline(1.0, color="black", linewidth=0.5, linestyle="--", alpha=0.5)
    ax.set_xlabel("Training step"); ax.set_ylabel("Edge / Flat densification-event ratio")
    ax.legend(loc="best", frameon=False)
    ax.set_title("Densification event redistribution over training")
    plt.tight_layout(); plt.savefig(args.out, dpi=120, bbox_inches="tight")
    print(f"saved {args.out}")


if __name__ == "__main__":
    main()