#!/usr/bin/env python3
"""Render mesh-swapped MimicGen demos into GR00T 1.5 LeRobot-style dataset.

Reads gen HDF5s from $WORKSPACE/gen_per_target_sam3d_50/<cat>_<model>/<cat>_<model>_gen/demo.hdf5,
writes parquet+meta+videos+extras to $DATASET_OUT.

This is a lightly-modified copy of /lp-dev/jonghoon/mimicgen_augment/render_per_target.py
that takes input/output paths from env vars (or defaults), so we can produce a fresh dataset
without clobbering the existing groot1.5_mimcgen_per_target.
"""
import os
os.environ["MUJOCO_GL"] = "egl"

import h5py, json, pathlib, glob, sys, gzip, re
import numpy as np
import pyarrow.parquet as pq
import pyarrow as pa

import robosuite, robocasa
from robomimic.envs.env_robosuite import EnvRobosuite
import robomimic.utils.obs_utils as ObsUtils
ObsUtils.initialize_obs_utils_with_obs_specs({'obs': {'low_dim': [], 'rgb': []}})

sys.path.insert(0, "/lp-dev/jonghoon/mimicgen_augment")
from swap_and_render import (
    extract_obs_state, write_meta, CAMERAS, IMG_SIZE,
    _bbox_diag, _bbox_diag_multi, _find_canonical_visual,
)
from build_sam3d_sources import find_sam3d_mesh as _find_sam3d_mesh_untextured

# The meshes build_sam3d_sources points at carry vertex colours but no UVs, so MuJoCo
# renders them flat grey. /data/minha/pose6daug/mesh_textured holds the same
# reconstructions re-baked with UVs + material_0.png, in the same <model>_to_* layout.
TEXTURED_PAIR_ROOT = pathlib.Path(os.environ.get(
    "TEXTURED_PAIR_ROOT", "/data/minha/pose6daug/mesh_textured/pairs"))


def find_sam3d_mesh(target_cat, target_model):
    for root in sorted(TEXTURED_PAIR_ROOT.glob("*")):
        for pair_dir in sorted(root.glob(f"{target_model}_to_*")):
            obj = pair_dir / "target_rotated" / "original.obj"
            if obj.exists():
                return obj
    return _find_sam3d_mesh_untextured(target_cat, target_model)
import mujoco, imageio

GEN_DIR = pathlib.Path(os.environ.get(
    "GEN_DIR", "/lp-dev/jonghoon/mimicgen_aug_groot15/gen_per_target_sam3d_50"))
DATASET_ROOT = pathlib.Path(os.environ.get(
    "DATASET_OUT", "/lp-dev/jonghoon/mimicgen_aug_groot15/dataset/groot1.5_mimcgen_per_target_50"))

TARGETS = [
    ("donut",          "donut_5"),
    ("jar",            "Jar023"),
    ("measuring_cup",  "MeasuringCup009"),
    ("soap_dispenser", "SoapDispenser010"),
    ("steak",          "steak_8"),
    ("syrup_bottle",   "SyrupBottle005"),
    ("teapot",         "teapot_6"),
    ("teapot",         "teapot_7"),
]


def swap_visual_to_sam3d(model_xml, sam3d_obj_path):
    tex = pathlib.Path(sam3d_obj_path).parent / "material_0.png"
    swap_texture = tex if tex.exists() else None
    canonical_file, canonical_scale_xyz = _find_canonical_visual(model_xml)
    if canonical_file is None:
        return None
    sam3d_diag = _bbox_diag(str(sam3d_obj_path))
    # _find_canonical_visual returns a LIST of visual submeshes (a robocasa object can be
    # split into several, e.g. teapot = lid + base) and its docstring says to take the union
    # bbox with _bbox_diag_multi. Upstream passed that list to _bbox_diag, which wants one
    # path, so the diagonal came back None and swap_visual_to_sam3d silently returned None --
    # the swap never happened, on any target.
    canonical_diag = (_bbox_diag_multi(canonical_file)
                      if isinstance(canonical_file, (list, tuple))
                      else _bbox_diag(canonical_file))
    if sam3d_diag is None or canonical_diag is None:
        return None
    mean_canonical = float(np.mean(canonical_scale_xyz))
    target_scale = (canonical_diag * mean_canonical) / sam3d_diag
    scale_str = f'scale="{target_scale:.6f} {target_scale:.6f} {target_scale:.6f}"'

    out = []
    for line in model_xml.split("\n"):
        s = line.strip()
        if s.startswith("<mesh") and 'name="obj_' in s and "collision" not in s.lower():
            line = re.sub(r'file="[^"]+"', f'file="{sam3d_obj_path}"', line)
            line = re.sub(r' scale="[^"]+"', f' {scale_str}', line)
            if "scale=" not in line:
                line = line.rstrip("/>").rstrip() + f" {scale_str}/>"
        elif s.startswith("<texture") and 'name="obj_' in s:
            if swap_texture is not None:
                # the replacement mesh is UV-mapped and ships its own baked texture
                line = re.sub(r'file="[^"]+"', f'file="{swap_texture}"', line)
                line = re.sub(r' builtin="[^"]+"', "", line)
            else:
                line = re.sub(r'file="[^"]+"', 'builtin="flat"', line)
                if "width=" not in line:
                    line = line.rstrip("/>").rstrip() + ' width="2" height="2"/>'
        elif s.startswith("<material") and 'name="obj_' in s:
            if swap_texture is None:
                # no texture to apply -- flat grey beats whatever the stale binding gives
                forced_rgba = "0.7 0.7 0.7 1.0"
                if re.search(r'rgba="[^"]+"', line):
                    line = re.sub(r'rgba="[^"]+"', f'rgba="{forced_rgba}"', line)
                else:
                    line = line.rstrip("/>").rstrip() + f' rgba="{forced_rgba}"/>'
                line = re.sub(r' texture="[^"]+"', "", line)
            elif re.search(r'rgba="[^"]+"', line):
                line = re.sub(r'rgba="[^"]+"', 'rgba="1 1 1 1"', line)
        out.append(line)
    return "\n".join(out)


def render_demo(env, demo_grp, out_root, episode_index, row_info, global_frame_offset, sam3d_obj_path=None):
    out_root = pathlib.Path(out_root)
    model_xml = demo_grp.attrs["model_file"]
    if sam3d_obj_path is not None:
        swapped = swap_visual_to_sam3d(model_xml, str(sam3d_obj_path))
        if swapped is not None:
            model_xml = swapped

    states = demo_grp["states"][0]
    state = dict(model=model_xml, states=states)
    if "ep_meta" in demo_grp.attrs:
        state["ep_meta"] = demo_grp.attrs["ep_meta"]

    try:
        env.reset_to(state)
    except Exception as e:
        print(f"  Reset failed: {e}")
        return False, 0

    all_states = demo_grp["states"][:]
    all_actions = demo_grp["actions"][:]
    sim = env.base_env.sim
    mj_model = sim.model._model
    mj_data = sim.data._data
    n_frames = len(all_states)

    cam_ids = {cam: mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_CAMERA, cam) for cam in CAMERAS}
    scene_option = mujoco.MjvOption()
    scene_option.geomgroup[:] = 0
    scene_option.geomgroup[1] = 1
    scene_option.sitegroup[:] = 0

    renderer = mujoco.Renderer(mj_model, height=IMG_SIZE, width=IMG_SIZE)
    ep_tag = f"episode_{episode_index:06d}"
    vid_dir = out_root / "videos" / "chunk-000"
    writers = {}
    for cam in CAMERAS:
        cam_dir = vid_dir / f"observation.images.{cam}"
        cam_dir.mkdir(parents=True, exist_ok=True)
        writers[cam] = imageio.get_writer(
            str(cam_dir / f"{ep_tag}.mp4"), fps=20, codec="libx264",
            macro_block_size=1, ffmpeg_params=["-pix_fmt", "yuv420p"],
        )

    obs_states = []
    for t in range(n_frames):
        sim.set_state_from_flattened(all_states[t])
        sim.forward()
        obs_states.append(extract_obs_state(env))
        for cam in CAMERAS:
            renderer.update_scene(mj_data, camera=cam_ids[cam], scene_option=scene_option)
            writers[cam].append_data(renderer.render())

    renderer.close()
    for w in writers.values():
        w.close()

    obs_states = np.array(obs_states)
    try:
        ep_meta_obj = json.loads(demo_grp.attrs.get("ep_meta", "{}"))
    except Exception:
        ep_meta_obj = {}

    # The HDF5 holds actions in the simulator's own order,
    #   [eef_pos(3), eef_rot(3), gripper, base(3), torso, base_mode]
    # while meta/modality.json declares the LeRobot order,
    #   [base_motion(4), control_mode, eef_pos(3), eef_rot(3), gripper_close].
    # Upstream wrote them through unchanged, so every slice landed on the wrong dimensions:
    # arm motion was read as base motion and base_mode as the gripper, which trains a policy
    # that drives the base away and never closes the gripper (measured 0/160 at eval).
    actions = all_actions[:, [7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5, 6]]
    if len(actions) < n_frames:
        pad = np.zeros((n_frames - len(actions), actions.shape[1]))
        actions = np.vstack([actions, pad])

    records = []
    for t in range(n_frames):
        is_last = (t == n_frames - 1)
        records.append({
            "observation.state": obs_states[t].tolist(),
            "action": actions[t].tolist(),
            "next.reward": 1.0 if is_last else 0.0,
            "next.done": is_last,
            "timestamp": t / 20.0,
            "frame_index": t,
            "episode_index": episode_index,
            "index": global_frame_offset + t,
            "task_index": 0,
            "annotation.human.task_description": 0,
            "annotation.human.task_name": 0,
        })

    schema = pa.schema([
        ("observation.state", pa.list_(pa.float64())),
        ("action", pa.list_(pa.float64())),
        ("next.reward", pa.float32()),
        ("next.done", pa.bool_()),
        ("timestamp", pa.float32()),
        ("frame_index", pa.int64()),
        ("episode_index", pa.int64()),
        ("index", pa.int64()),
        ("task_index", pa.int64()),
        ("annotation.human.task_description", pa.int64()),
        ("annotation.human.task_name", pa.int64()),
    ])
    table = pa.Table.from_pylist(records, schema=schema)
    data_dir = out_root / "data" / "chunk-000"
    data_dir.mkdir(parents=True, exist_ok=True)
    pq.write_table(table, str(data_dir / f"{ep_tag}.parquet"))

    extras_dir = out_root / "extras" / ep_tag
    extras_dir.mkdir(parents=True, exist_ok=True)
    aug_meta = dict(ep_meta_obj)
    aug_meta["augmentation"] = {
        "type": "mimicgen_per_target_sam3d_50",
        "target_mesh": row_info["target_mesh"],
        "source_demo": row_info["source_demo"],
    }
    (extras_dir / "ep_meta.json").write_text(json.dumps(aug_meta, indent=2))
    with gzip.open(str(extras_dir / "model.xml.gz"), "wt") as f:
        f.write(model_xml)
    np.savez_compressed(str(extras_dir / "states.npz"), states=all_states)

    return True, n_frames


def main():
    DATASET_ROOT.mkdir(parents=True, exist_ok=True)
    print(f"GEN_DIR={GEN_DIR}")
    print(f"DATASET_ROOT={DATASET_ROOT}")

    target_demos = {}
    for cat, model in TARGETS:
        name = f"{cat}_{model}"
        gen_paths = sorted(glob.glob(f"{GEN_DIR}/{name}*/{name}_gen/demo.hdf5"))
        if not gen_paths:
            print(f"  WARN: no demos for {cat}/{model} in {GEN_DIR}/{name}")
            target_demos[(cat, model)] = []
            continue
        target_demos[(cat, model)] = gen_paths

    first_hdf5 = None
    for paths in target_demos.values():
        if paths:
            first_hdf5 = paths[0]
            break
    if first_hdf5 is None:
        print("No demos found at all!")
        return

    fh0 = h5py.File(first_hdf5, "r")
    env_meta = json.loads(fh0["data"].attrs["env_args"])
    fh0.close()
    env_kwargs = dict(env_meta["env_kwargs"])
    env_kwargs["has_offscreen_renderer"] = True
    env_kwargs["use_camera_obs"] = False
    env = EnvRobosuite(env_name=env_meta["env_name"], **env_kwargs)

    episodes_info = []
    task_descriptions = set()
    global_frame_offset = 0
    success_count = 0
    ep_idx = 0

    only = os.environ.get("TARGET_FILTER")   # "<cat>/<model>" -> shard this run to one target
    for cat, model in TARGETS:
        if only and f"{cat}/{model}" != only:
            continue
        gen_paths = target_demos[(cat, model)]
        target_label = f"{cat}/{model}"
        # The visual-mesh swap is OFF here. Upstream replaces the object's visual mesh with
        # a reconstruction from an eval rollout and forces it grey; neither the Ours nor the
        # VACE dataset builder does that, and our episodes were generated (and already
        # inspected in output_aug_256) with the plain RoboCasa asset. USE_MESH_SWAP=1 restores
        # the upstream behaviour.
        sam3d_mesh = find_sam3d_mesh(cat, model) if os.environ.get("USE_MESH_SWAP") == "1" else None
        print(f"  -> {target_label}: mesh swap = {sam3d_mesh}")
        n_for_target = 0
        for hp in gen_paths:
            fh = h5py.File(hp, "r")
            for dk in sorted(fh["data"].keys(), key=lambda x: int(x.split("_")[1])):
                row_info = {
                    "target_mesh": target_label,
                    "source_demo": f"{hp}::{dk}",
                }

                ep_tag = f"episode_{ep_idx:06d}"
                parquet_path = DATASET_ROOT / "data" / "chunk-000" / f"{ep_tag}.parquet"
                if parquet_path.exists():
                    tbl = pq.read_table(str(parquet_path))
                    n = len(tbl)
                    em_path = DATASET_ROOT / "extras" / ep_tag / "ep_meta.json"
                    td = "Pick and place"
                    if em_path.exists():
                        em = json.loads(em_path.read_text())
                        td = em.get("lang", td)
                    episodes_info.append({"episode_index": ep_idx, "tasks": [td], "length": n})
                    task_descriptions.add(td)
                    global_frame_offset += n
                    success_count += 1
                    n_for_target += 1
                    ep_idx += 1
                    print(f"  ep{ep_idx-1:03d} [{target_label}] {dk} — already done")
                    continue

                print(f"  ep{ep_idx:03d} [{target_label}] {dk}", flush=True)
                ok, n_frames = render_demo(
                    env, fh[f"data/{dk}"],
                    DATASET_ROOT, ep_idx, row_info, global_frame_offset,
                    sam3d_obj_path=sam3d_mesh,
                )
                if ok:
                    success_count += 1
                    n_for_target += 1
                    em_path = DATASET_ROOT / "extras" / ep_tag / "ep_meta.json"
                    em = json.loads(em_path.read_text()) if em_path.exists() else {}
                    td = em.get("lang", f"Pick and place {cat}")
                    episodes_info.append({"episode_index": ep_idx, "tasks": [td], "length": n_frames})
                    task_descriptions.add(td)
                    global_frame_offset += n_frames
                    print(f"    OK ({n_frames} frames)")
                else:
                    print(f"    FAILED")
                ep_idx += 1
            fh.close()
        print(f"  -> {target_label}: {n_for_target} episodes")

    env.close()

    if episodes_info:
        write_meta(DATASET_ROOT, episodes_info, sorted(task_descriptions))
        print(f"Wrote meta/ with {len(episodes_info)} episodes")
    print(f"Done: {success_count} episodes total")


if __name__ == "__main__":
    main()