#!/usr/bin/env python3
"""Pick the first 256 episodes in true generation order and lay out the work for both
downstream stages.

Every worker writes each success to tmp/date_..._time_HH_MM_SS.hdf5 before merge_all_hdf5
sorts those files by timestamp into demo_0..demo_N and deletes the folder. snapshot_times.sh
captured the filenames while they existed, so joining them back gives each merged demo its
wall-clock time -- and since all 32 workers share one machine clock, a global ordering.

Writes:
  selection.json     the ordered 256 with their timestamps
  render_plan.json   for render_from_plan.py (preview mp4s, original asset)
  gen/shard_K/...    per-shard gen trees for render_to_gr00t.py (dataset, swapped mesh)
"""
import json, os, re, glob, sys, collections, datetime
import h5py

S = "/tmp/claude-1001/-home-nvidia/68c27ffb-706b-4cb3-8bf2-4618a200d378/scratchpad"
NAT = f"{S}/natural256"
SRC = f"{S}/big/src150_fix.hdf5"
N_KEEP = int(os.environ.get("N_KEEP", 256))
N_CONV_SHARDS = int(os.environ.get("N_CONV_SHARDS", 24))

CATEGORY = {
    "donut_5": "donut", "Jar023": "jar", "MeasuringCup009": "measuring_cup",
    "SoapDispenser010": "soap_dispenser", "steak_8": "steak",
    "SyrupBottle005": "syrup_bottle", "teapot_6": "teapot", "teapot_7": "teapot",
}

# ---- source demo index -> (demo key, the object that source episode had) ----
src = {}
with h5py.File(SRC, "r") as f:
    keys = [k.decode() for k in f["mask/valid"][:]]
    for i, k in enumerate(keys):
        meta = json.loads(f[f"data/{k}"].attrs["ep_meta"])
        cat = next((c["info"]["cat"] for c in meta.get("object_cfgs", [])
                    if c.get("name") == "obj"), "unknown")
        src[i] = (k, cat)

# ---- timestamps captured before the tmp folders were deleted ----
# One line per poll, so the same file appears many times and its mtime drifts while it is
# still being written -- dedupe on the name and order by the timestamp the name itself
# carries, which is fixed at creation.
seen = collections.defaultdict(dict)
with open(f"{NAT}/episode_times.tsv") as f:
    for line in f:
        parts = line.rstrip("\n").split("\t")
        if len(parts) != 3:
            continue
        tmpdir, fname, _ = parts
        m = re.match(r"date_(\d+)_(\d+)_(\d+)_time_(\d+)_(\d+)_(\d+)\.hdf5", fname)
        if not m:
            continue
        worker = os.path.basename(os.path.dirname(os.path.dirname(tmpdir)))
        mo, dy, yr, hh, mm, ss = (int(x) for x in m.groups())
        seen[worker][fname] = datetime.datetime(
            yr, mo, dy, hh, mm, ss, tzinfo=datetime.timezone.utc).timestamp()
times = {w: sorted((v, k) for k, v in d.items()) for w, d in seen.items()}

rows, missing = [], []
for wdir in sorted(glob.glob(f"{NAT}/nat_*")) + sorted(glob.glob(f"{NAT}/r2_*")):
    worker = os.path.basename(wdir)
    model = re.sub(r"^(nat|r2)_|_w\d+$", "", worker)
    hp = f"{wdir}/{worker}/demo.hdf5"
    if not os.path.isdir(wdir) or not os.path.exists(hp):
        continue
    with h5py.File(hp, "r") as f:
        demos = sorted(f["data"].keys(), key=lambda x: int(x.split("_")[1]))
        stamps = list(times.get(worker, []))
        # The poller can miss the very last success: the worker merges and deletes tmp/
        # immediately after writing it. Only the tail can ever be missing (earlier files sat
        # there for the whole run), so pad it with the merged file's own mtime -- the moment
        # the worker finished, which bounds that episode from above.
        if len(stamps) < len(demos):
            end = os.path.getmtime(hp)
            pad = len(demos) - len(stamps)
            missing.append(f"{worker}: padded {pad} with run end")
            stamps += [(end, "<run end>")] * pad
        for i, dk in enumerate(demos):
            g = f[f"data/{dk}"]
            si = int(g["src_demo_inds"][0]) if "src_demo_inds" in g else -1
            sk, scat = src.get(si, (f"idx{si}", "unknown"))
            t = stamps[i][0] if i < len(stamps) else float("inf")
            rows.append(dict(t=t, worker=worker, run_dir=f"{wdir}/{worker}", demo=dk,
                             target=model, source_demo=sk, source_object=scat))

rows.sort(key=lambda r: (r["t"], r["worker"], r["demo"]))
print(f"총 성공 {len(rows)}개" + (f"  (경고: {missing})" if missing else ""))
for i, r in enumerate(rows):
    r["order"] = i
kept = rows[:N_KEEP]

taken = collections.Counter()
for r in kept:
    base = f"{r['source_demo'].replace('demo_', 'ep')}_{r['source_object']}_{r['target']}"
    taken[base] += 1
    r["name"] = base if taken[base] == 1 else f"{base}_{taken[base]}"

json.dump(kept, open(f"{NAT}/selection.json", "w"), indent=1)
json.dump([{k: r[k] for k in ("run_dir", "demo", "target", "name")} for r in kept],
          open(f"{NAT}/render_plan.json", "w"), indent=1)

# ---- conversion shards: contiguous chunks, each its own GEN_DIR tree ----
by_target = collections.defaultdict(list)
for r in kept:
    by_target[r["target"]].append(r)
chunks = []
per = max(1, -(-len(kept) // N_CONV_SHARDS))
for model, rs in by_target.items():
    for i in range(0, len(rs), per):
        chunks.append((model, rs[i:i + per]))

for k, (model, rs) in enumerate(chunks):
    name = f"{CATEGORY[model]}_{model}"
    d = f"{NAT}/gen/shard_{k:02d}/{name}/{name}_gen"
    os.makedirs(d, exist_ok=True)
    out = f"{d}/demo.hdf5"
    if os.path.exists(out):
        os.remove(out)
    with h5py.File(out, "w") as fo:
        grp = fo.create_group("data")
        total, env_args = 0, None
        by_run = collections.defaultdict(list)
        for r in rs:
            by_run[r["run_dir"]].append(r)
        j = 0
        for run_dir, items in by_run.items():
            with h5py.File(f"{run_dir}/demo.hdf5", "r") as fi:
                if env_args is None:
                    env_args = fi["data"].attrs["env_args"]
                for r in items:
                    fi.copy(f"data/{r['demo']}", grp, name=f"demo_{j}")
                    total += int(fi[f"data/{r['demo']}"].attrs["num_samples"])
                    j += 1
        grp.attrs["total"] = total
        grp.attrs["env_args"] = env_args

print(f"선택 {len(kept)}개 -> render_plan.json, 변환 샤드 {len(chunks)}개 "
      f"(샤드당 최대 {per})")
c = collections.Counter(r["target"] for r in kept)
print("\n물체별 분포:")
for k, v in c.most_common():
    print(f"  {k:18s} {v}")