import mujoco
import numpy as np
import pytest
from conftest import GOOD_SCENE

pytest.importorskip("pxr")
from pxr import Usd, UsdGeom  # noqa: E402

from api.mjcf_check import prepare_scene_dir  # noqa: E402
from worker.usd_scene import build_stage  # noqa: E402


def test_replay_stage_matches_mujoco(tmp_path, menagerie):
    """Every drawn geom's world position in the USD equals MuJoCo's geom_xpos, frame by frame."""
    scene_dir = prepare_scene_dir(GOOD_SCENE, "unitree_g1", menagerie)
    model = mujoco.MjModel.from_xml_path(str(scene_dir / "scene.xml"))
    data = mujoco.MjData(model)
    data.qvel[:] = np.random.default_rng(0).normal(0, 1, model.nv)  # tumble things around
    pos, quat, geom_xpos = [], [], []
    for _ in range(20):
        for _ in range(10):
            mujoco.mj_step(model, data)
        pos.append(data.xpos.copy())
        quat.append(data.xquat.copy())
        geom_xpos.append(data.geom_xpos.copy())
    names = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, b) or f"body_{b}" for b in range(model.nbody)]
    np.savez(tmp_path / "traj.npz", body_pos=pos, body_quat=quat, body_names=names, fps=30)

    info = build_stage(scene_dir / "scene.xml", tmp_path / "traj.npz", tmp_path / "scene.usdc", "pelvis")
    assert info["frames"] == 20 and info["geoms_drawn"] >= 3
    stage = Usd.Stage.Open(str(tmp_path / "scene.usdc"))
    checked = 0
    for prim in stage.Traverse():
        name = prim.GetName()
        if not (name.startswith("g") and name[1:].isdigit()):
            continue
        g = int(name[1:])
        if model.geom_type[g] == mujoco.mjtGeom.mjGEOM_PLANE:
            continue
        for i in (0, 9, 19):
            m = UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(Usd.TimeCode(i))
            centre = np.array(m.Transform((0, 0, 0))) if model.geom_type[g] != mujoco.mjtGeom.mjGEOM_MESH else None
            if centre is not None:
                assert np.allclose(centre, geom_xpos[i][g], atol=1e-4), (g, i)
                checked += 1
    assert checked >= 6
    cam = UsdGeom.Xformable(stage.GetPrimAtPath("/World/Camera"))
    eye = np.array(cam.ComputeLocalToWorldTransform(Usd.TimeCode(0)).ExtractTranslation())
    assert eye[2] > pos[0][1][2]  # looking down at the robot