"""Cache builder for OGBench play datasets (scene, puzzle-3x3).

Play trajectories do not terminate at a goal, so the goal for each transition is
a hindsight future observation sampled within the same trajectory (fixed seed).
Nested budgets use the same fixed-permutation-prefix scheme as the cube release.
"""
import argparse, json, os
import numpy as np, ogbench
from common import get_traj_bounds, nested_traj_order

ap = argparse.ArgumentParser()
ap.add_argument("--dataset", required=True)          # e.g. scene-play-v0
ap.add_argument("--budget", type=int, required=True) # trajectories
ap.add_argument("--split_seed", type=int, default=0)
ap.add_argument("--goal_seed", type=int, default=0)
ap.add_argument("--outdir", default="cache")
a = ap.parse_args()

_, train_ds, _ = ogbench.make_env_and_datasets(a.dataset, compact_dataset=False)
obs = train_ds["observations"].astype(np.float32)
act = train_ds["actions"].astype(np.float32)
term = train_ds["terminals"].astype(np.float32)
bounds = get_traj_bounds(term)
order = nested_traj_order(len(bounds), seed=a.split_seed)
sel = np.sort(order[:a.budget])

rng = np.random.default_rng(a.goal_seed)
O, A, G = [], [], []
for ti in sel:
    s, e = bounds[int(ti)]
    idx = np.arange(s, e)
    # hindsight future goal: uniform index in [t, e-1] for each t
    u = rng.random(len(idx))
    gj = (idx + (u * (e - idx)).astype(np.int64)).clip(max=e - 1)
    O.append(obs[idx]); A.append(act[idx]); G.append(obs[gj])
O = np.concatenate(O); A = np.concatenate(A); G = np.concatenate(G)
os.makedirs(a.outdir, exist_ok=True)
tag = a.dataset.replace("-play-v0", "")
out = f"{a.outdir}/{tag}_state_b{a.budget:04d}.npz"
np.savez(out, observation=O, action=A, goal=G)
print(json.dumps({"dataset": a.dataset, "budget": a.budget, "trajectories": len(sel),
                  "transitions": int(len(O)), "obs_dim": int(O.shape[1]),
                  "act_dim": int(A.shape[1]), "cache": out}))