#!/usr/bin/env python3
"""qwen_dataset_gen.py — self-labelling grounding data from the RX1 MuJoCo
scene (pipeline edition, no external repo dependency).

Randomizes cube positions, renders the ego camera, and projects each cube's
true world position into the image to produce Qwen-VLA grounding rows:

    grounding.jsonl   {"image", "instruction", "target", "pixel": [u, v], "in_view"}
    images/frame_<n>.png

Same schema as world_model_test/rx1_brain/dataset_gen.py, generated from the
zelpi sim instead of the PyBullet twin. Last line: PIPELINE_RESULT {...}.

    MUJOCO_GL=glfw python qwen_dataset_gen.py --episodes 40 --out <dir>
"""
from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

import mujoco  # noqa: E402
from PIL import Image  # noqa: E402

import scripts.stack_demo as sd  # noqa: E402
from env.rx1_env import RX1Env  # noqa: E402

IMG = 512
CAMERA = "ego"
COLORS = ("red", "blue", "green", "yellow")
SPAWN_X = (0.20, 0.50)
SPAWN_Y = (-0.30, 0.12)
TABLE_Z = 0.445
MIN_SEP = 0.08
TEMPLATES = (
    "Where is the {c} cube? Answer with its pixel coordinates as [u, v].",
    "Point to the {c} cube. Reply only with [u, v] pixel coordinates.",
    "Locate the {c} cube in the image and answer [u, v].",
)


def emit(tag: str, payload: dict):
    print(f"{tag} {json.dumps(payload)}", flush=True)


def set_cube(env, cube, pos):
    j = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, f"{cube}_free")
    a = env.model.jnt_qposadr[j]
    env.data.qpos[a:a + 3] = pos
    env.data.qpos[a + 3:a + 7] = [1, 0, 0, 0]
    env.data.qvel[env.model.jnt_dofadr[j]:env.model.jnt_dofadr[j] + 6] = 0


def project(env, cam_id, world, w, h):
    """World point → pixel in a MuJoCo fixed/tracked camera (looks along -z)."""
    pos = env.data.cam_xpos[cam_id]
    rot = env.data.cam_xmat[cam_id].reshape(3, 3)
    p = rot.T @ (np.asarray(world) - pos)
    depth = -p[2]
    if depth <= 1e-6:
        return None
    f = (h / 2) / math.tan(math.radians(env.model.cam_fovy[cam_id]) / 2)
    u = w / 2 + f * p[0] / depth
    v = h / 2 - f * p[1] / depth
    return float(u), float(v), float(depth)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--episodes", type=int, default=40, help="scene randomizations")
    ap.add_argument("--out", required=True)
    ap.add_argument("--seed", type=int, default=11)
    a = ap.parse_args()

    out = Path(a.out)
    (out / "images").mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(a.seed)

    env = RX1Env(sd.ROBOT_CFG)
    cam_id = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_CAMERA, CAMERA)
    r = mujoco.Renderer(env.model, IMG, IMG)

    mode = "a" if (out / "grounding.jsonl").exists() else "w"
    n_rows = 0
    with open(out / "grounding.jsonl", mode) as fh:
        start = len(list((out / "images").glob("frame_*.png")))
        for ep in range(a.episodes):
            env.reset()
            spots = []
            for _ in range(100):
                p = np.array([rng.uniform(*SPAWN_X), rng.uniform(*SPAWN_Y)])
                if all(np.linalg.norm(p - q) >= MIN_SEP for q in spots):
                    spots.append(p)
                if len(spots) == len(COLORS):
                    break
            if len(spots) < len(COLORS):
                continue
            for color, xy in zip(COLORS, spots):
                set_cube(env, f"{color}_cube", [xy[0], xy[1], TABLE_Z])
            mujoco.mj_forward(env.model, env.data)

            r.update_scene(env.data, camera=CAMERA)
            frame = r.render()
            img_name = f"images/frame_{start + ep:05d}.png"
            Image.fromarray(frame).save(out / img_name)

            for color in COLORS:
                body = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY, f"{color}_cube")
                world = env.data.xpos[body]
                proj = project(env, cam_id, world, IMG, IMG)
                in_view = proj is not None and 0 <= proj[0] < IMG and 0 <= proj[1] < IMG
                row = {
                    "image": img_name,
                    "instruction": TEMPLATES[int(rng.integers(len(TEMPLATES)))].format(c=color),
                    "target": f"{color}_cube",
                    "pixel": [round(proj[0], 1), round(proj[1], 1)] if in_view else None,
                    "in_view": bool(in_view),
                }
                fh.write(json.dumps(row) + "\n")
                n_rows += 1
            if (ep + 1) % 10 == 0:
                emit("PIPELINE_PROGRESS", {"step": ep + 1, "loss": None,
                                           "note": f"{n_rows} rows"})

    emit("PIPELINE_RESULT", {"ok": True, "rows": n_rows, "out": str(out)})


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        emit("PIPELINE_RESULT", {"ok": False, "message": f"{type(e).__name__}: {e}"})
        raise