#!/usr/bin/env python3
"""Render episodes listed in a render plan, one mp4 per camera.

The plan fixes every folder name up front, in one global pass over all runs, so a name is
unique across the whole output rather than only within the run that produced it. Assigning
names per run is what silently dropped episodes: a later run's success collided with an
earlier one's folder and was skipped as already done.

Plan entries are `{run_dir, demo, target, name}`. Shard with `--shard i --num-shards n` to
render in parallel; each episode writes a `.done` marker, so re-running skips finished work.

    MUJOCO_GL=egl python render_from_plan.py --plan render_plan.json \
        --source src150_fix.hdf5 --filter-key valid --out output_aug_256 --shard 0 --num-shards 6
"""
from __future__ import annotations

import argparse
import json
import os
import time

os.environ.setdefault("MUJOCO_GL", "egl")
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")

import h5py
import imageio
import mujoco
import numpy as np
import robocasa  # noqa: F401
import robosuite
from robosuite.controllers import load_composite_controller_config

CAMERAS = ["robot0_agentview_left", "robot0_agentview_right", "robot0_eye_in_hand"]


def source_index(path, filter_key):
    out = {}
    with h5py.File(path, "r") as f:
        keys = ([k.decode() for k in f[f"mask/{filter_key}"][:]] if filter_key
                else sorted(f["data"].keys(), key=lambda x: int(x.split("_")[1])))
        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")
            out[i] = (k, cat)
    return out


def build_env(env_args, res):
    kw = dict(env_args["env_kwargs"])
    kw.pop("env_name", None)
    kw.update(has_renderer=False, has_offscreen_renderer=True, use_camera_obs=False,
              use_object_obs=True, camera_depths=False, camera_names=CAMERAS,
              camera_widths=res, camera_heights=res, ignore_done=True)
    kw.setdefault("controller_configs",
                  load_composite_controller_config(controller=None, robot="PandaOmron"))
    return robosuite.make(env_name=env_args["env_name"], **kw)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--plan", required=True)
    ap.add_argument("--source", required=True)
    ap.add_argument("--filter-key", default=None)
    ap.add_argument("--out", required=True)
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--num-shards", type=int, default=1)
    ap.add_argument("--res", type=int, default=256)
    ap.add_argument("--fps", type=int, default=20)
    a = ap.parse_args()

    plan = [e for i, e in enumerate(json.load(open(a.plan))) if i % a.num_shards == a.shard]
    src = source_index(a.source, a.filter_key)
    os.makedirs(a.out, exist_ok=True)

    env, opened, done = None, None, 0
    for e in plan:
        ep_dir = os.path.join(a.out, e["name"])
        if os.path.exists(os.path.join(ep_dir, ".done")):
            done += 1
            continue
        t0 = time.time()
        with h5py.File(os.path.join(e["run_dir"], "demo.hdf5"), "r") as f:
            g = f[f"data/{e['demo']}"]
            states = g["states"][:]
            meta = json.loads(g.attrs["ep_meta"])
            model_xml = g.attrs["model_file"]
            si = int(g["src_demo_inds"][0]) if "src_demo_inds" in g else -1
            env_args = json.loads(f["data"].attrs["env_args"])
        sk, scat = src.get(si, (f"idx{si}", "unknown"))

        if env is None:
            env = build_env(env_args, a.res)
        env.set_ep_meta(meta)
        env.reset()
        env.reset_from_xml_string(model_xml)
        env.sim.reset()

        model, data = env.sim.model._model, env.sim.data._data
        opt = mujoco.MjvOption()
        opt.geomgroup[:] = 0
        opt.geomgroup[1] = 1          # visual geoms, as the dataset's own videos are drawn
        opt.sitegroup[:] = 0
        r = mujoco.Renderer(model, height=a.res, width=a.res)
        cam_ids = {c: mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, c) for c in CAMERAS}
        frames = {c: [] for c in CAMERAS}
        for t in range(len(states)):
            env.sim.set_state_from_flattened(states[t])
            env.sim.forward()
            for c in CAMERAS:
                r.update_scene(data, camera=cam_ids[c], scene_option=opt)
                frames[c].append(r.render().copy())
        r.close()

        for c in CAMERAS:
            d = os.path.join(ep_dir, c)
            os.makedirs(d, exist_ok=True)
            imageio.mimwrite(os.path.join(d, f"{c}.mp4"), frames[c], fps=a.fps,
                             codec="libx264", macro_block_size=1,
                             ffmpeg_params=["-pix_fmt", "yuv420p"])
        with open(os.path.join(ep_dir, ".done"), "w") as fh:
            json.dump(dict(generated_demo=e["demo"], run_dir=e["run_dir"],
                           source_demo=sk, source_object=scat, target_object=e["target"],
                           frames=len(states), layout=meta.get("layout_id"),
                           style=meta.get("style_id"), lang=meta.get("lang")), fh, indent=1)
        done += 1
        print(f"  [{a.shard}] {e['name']}: {len(states)}f {time.time()-t0:.0f}s", flush=True)
    print(f"shard {a.shard}: {done}/{len(plan)}", flush=True)


if __name__ == "__main__":
    raise SystemExit(main())