"""Re-render dataset trajectories INSIDE the OGBench eval env for domain-matched pixels.

The release stores full sim state (qpos/qvel); we replay it into the evaluation
environment and call get_pixel_observation(), so training images come from the
same renderer/camera as evaluation. Also stores the per-episode goal frame
(final state) and the 28-d state obs for the proprioception ablation.
"""
import argparse, json, time
from pathlib import Path
import numpy as np, gymnasium as gym, ogbench, mujoco

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="data/bucket/ogbench-cube-success-1000-v1")
    ap.add_argument("--budget", type=int, required=True)
    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(parents=True, exist_ok=True)

    env = gym.make(a.env); u = env.unwrapped
    env.reset(seed=0, options=dict(task_id=1, render_goal=False))
    probe = np.asarray(u.get_pixel_observation())
    H, W, C = probe.shape
    print(f"[render] env={a.env} frame={probe.shape} dtype={probe.dtype}", flush=True)

    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)
    frames = np.zeros((total, H, W, C), np.uint8)
    goals = np.zeros((total, H, W, C), np.uint8)
    acts = np.zeros((total, 5), np.float32)
    states = np.zeros((total, 28), 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 = z["qpos"], z["qvel"]
            act, sobs = z["action"], z["observation"]
        T = len(act) - 1
        # goal frame = final (terminal) state of this successful demo
        u.data.qpos[:] = qpos[-1]; u.data.qvel[:] = qvel[-1]
        mujoco.mj_forward(u.model, u.data)
        gframe = np.asarray(u.get_pixel_observation()).copy()
        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()
            goals[i] = gframe
            acts[i] = act[t]; states[i] = sobs[t]; tids[i] = ti
            i += 1
        if (ti + 1) % 100 == 0:
            print(f"  {ti+1}/{len(rows)} episodes, {i} frames, {i/(time.time()-t0):.0f} fps", flush=True)
    assert i == total, (i, total)
    tag = f"b{a.budget:04d}"
    np.save(out / f"pix_{tag}_frames.npy", frames)
    np.save(out / f"pix_{tag}_goals.npy", goals)
    np.savez(out / f"pix_{tag}_meta.npz", action=acts, state=states, traj_id=tids)
    print(json.dumps({"budget": a.budget, "transitions": total, "frame": [H, W, C],
                      "seconds": round(time.time() - t0, 1),
                      "frames_gb": round(frames.nbytes / 1e9, 3)}), flush=True)

if __name__ == "__main__":
    main()