"""Re-render the release at 224x224 inside the eval env (OpenVLA's native size).

Raises the model's offscreen framebuffer (hard-set to 64) before swapping in a
224 renderer. Stores frames as a memmap plus per-transition action and goal xy.
"""
import argparse, json, time
from pathlib import Path
import numpy as np, gymnasium as gym, ogbench, mujoco

ap = argparse.ArgumentParser()
ap.add_argument("--root", default="data/bucket/ogbench-cube-success-1000-v1")
ap.add_argument("--budget", type=int, default=1000)
ap.add_argument("--env", default="visual-cube-single-v0")
ap.add_argument("--outdir", default="cache")
a = ap.parse_args()
root = Path(a.root); out = Path(a.outdir); out.mkdir(exist_ok=True)

env = gym.make(a.env); u = env.unwrapped
env.reset(seed=0, options=dict(task_id=1, render_goal=False))
u.model.vis.global_.offwidth = 224
u.model.vis.global_.offheight = 224
u._renderer = mujoco.Renderer(u.model, height=224, width=224)
u._render_height = 224; u._render_width = 224
assert np.asarray(u.get_pixel_observation()).shape[:2] == (224, 224)

rows = [json.loads(l) for l in (root / "subsets" / f"subset-{a.budget:04d}.jsonl").read_text().splitlines()]
total = sum(r["transition_count"] for r in rows)
tag = f"b{a.budget:04d}"
frames = np.lib.format.open_memmap(out / f"vla_{tag}_frames224.npy", mode="w+",
                                   dtype=np.uint8, shape=(total, 224, 224, 3))
acts = np.zeros((total, 5), np.float32)
goals = np.zeros((total, 2), np.float32)
tids = np.zeros(total, np.int32)

i, t0 = 0, time.time()
for ti, row in enumerate(rows):
    with np.load(root / "data" / row["archive_file"], allow_pickle=False) as z:
        qpos, qvel, act = z["qpos"], z["qvel"], z["action"]
        gxy = np.asarray(z["variation.cube.goal_position"])[0, 0, :]   # (2,) world xy
    T = len(act) - 1
    for t in range(T):
        u.data.qpos[:] = qpos[t]; u.data.qvel[:] = qvel[t]
        mujoco.mj_forward(u.model, u.data)
        frames[i] = u.get_pixel_observation()
        acts[i] = act[t]; goals[i] = gxy; tids[i] = ti
        i += 1
    if (ti + 1) % 200 == 0:
        print(f"  {ti+1}/{len(rows)} eps, {i} frames, {i/(time.time()-t0):.0f} fps", flush=True)
assert i == total
frames.flush()
np.savez(out / f"vla_{tag}_meta.npz", action=acts, goal_xy=goals, traj_id=tids)
print(json.dumps({"budget": a.budget, "transitions": total, "frame": [224, 224, 3],
                  "gb": round(frames.nbytes / 1e9, 2), "seconds": round(time.time() - t0, 1)}))