"""Build a compact training cache from the frozen ogbench-cube-success release.

Reads ONLY the small state fields (observation/action/success) per episode in
frozen subset order, skips the terminal NaN-action row, and uses each episode's
final observation as the (hindsight) goal for goal-conditioned BC.
"""
import argparse, json, os
from pathlib import Path
import numpy as np

BUDGETS = (10, 50, 100, 1000)

def build(root, budget, out):
    root = Path(root)
    rows = [json.loads(l) for l in (root / "subsets" / f"subset-{budget:04d}.jsonl").read_text().splitlines()]
    assert len(rows) == budget, f"expected {budget} rows, got {len(rows)}"
    obs_l, act_l, goal_l, succ_l, tid_l = [], [], [], [], []
    for ti, row in enumerate(rows):
        with np.load(root / "data" / row["archive_file"], allow_pickle=False) as z:
            o = z["observation"].astype(np.float32)   # (T+1, 28)
            a = z["action"].astype(np.float32)         # (T+1, 5), last row NaN
            s = z["success"].astype(np.bool_)          # (T+1,)
        T = len(a) - 1                                 # executed transitions
        assert np.isfinite(a[:T]).all(), "unexpected NaN in executed actions"
        goal = o[-1]                                   # final/achieved state = goal
        obs_l.append(o[:T]); act_l.append(a[:T])
        goal_l.append(np.broadcast_to(goal, (T, o.shape[1])))
        succ_l.append(s[:T]); tid_l.append(np.full(T, ti, np.int32))
    OBS = np.concatenate(obs_l); ACT = np.concatenate(act_l)
    GOAL = np.concatenate(goal_l); SUC = np.concatenate(succ_l); TID = np.concatenate(tid_l)
    os.makedirs(Path(out).parent, exist_ok=True)
    np.savez(out, observation=OBS, action=ACT, goal=GOAL, success=SUC, traj_id=TID)
    print(json.dumps({"budget": budget, "episodes": budget, "transitions": int(len(OBS)),
                      "obs_dim": int(OBS.shape[1]), "act_dim": int(ACT.shape[1]),
                      "success_transitions": int(SUC.sum()), "cache": out}))

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="data/bucket/ogbench-cube-success-1000-v1")
    ap.add_argument("--budget", type=int, choices=BUDGETS, default=1000)
    ap.add_argument("--out", default=None)
    a = ap.parse_args()
    build(a.root, a.budget, a.out or f"cache/cube_state_b{a.budget:04d}.npz")