"""Export YAM right-arm per-link meshes (in link-local frame) to OBJ for the web viewer."""
import argparse, sys, os
from isaaclab.app import AppLauncher
p=argparse.ArgumentParser(); AppLauncher.add_app_launcher_args(p); a=p.parse_args(); a.headless=True; a.enable_cameras=False
app=AppLauncher(a).app
import numpy as np, gymnasium as gym
REPO=os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0,os.path.join(REPO,"source"))
import bimanual.tasks.manager_based.yam  # noqa
from isaaclab_tasks.utils import parse_env_cfg
TASK="Template-YAM-Play-v0"
env=gym.make(TASK, cfg=parse_env_cfg(TASK, device="cuda:0", num_envs=1)); u=env.unwrapped; env.reset()
R=u.scene["right_robot"]
import omni.usd
from pxr import Usd, UsdGeom, Gf
stage=omni.usd.get_context().get_stage()
root_path=R.root_physx_view.prim_paths[0]
print("[m] root:", root_path, flush=True)
xf=UsdGeom.XformCache(Usd.TimeCode.Default())
OUT="/home/horde/project/xiaotong/outputs/yam_web/mesh"
os.makedirs(OUT, exist_ok=True)

def find_link_prim(name):
    for pr in Usd.PrimRange(stage.GetPrimAtPath(root_path)):
        if pr.GetName()==name:
            return pr
    return None

def gfmat_to_np(m):
    return np.array([[m[i][j] for j in range(4)] for i in range(4)])

LINKS=["arm","link_1","link_2","link_3","link_4","link_5","link_6","left_finger","right_finger"]
linkpath={}; Lwinv={}
for n in LINKS:
    lp=find_link_prim(n)
    if lp is None: continue
    linkpath[lp.GetPath().pathString]=n
    Lwinv[n]=np.linalg.inv(gfmat_to_np(xf.GetLocalToWorldTransform(lp)))
# fingers are rigidly (approx) under link_6; export them in link_6's frame so they
# mount onto link_6's group in the viewer (their small prismatic slide is ignored for viz).
FRAME={n:n for n in LINKS}
FRAME["left_finger"]="link_6"; FRAME["right_finger"]="link_6"
def owning_link(prim):
    p=prim.GetPath()
    while not p.isEmpty and p.pathString!="/":
        if p.pathString in linkpath: return linkpath[p.pathString]
        p=p.GetParentPath()
    return None
buck={n:{"V":[],"F":[]} for n in LINKS}
for pr in Usd.PrimRange(stage.GetPrimAtPath(root_path)):
    if pr.GetTypeName()!="Mesh": continue
    ln=owning_link(pr)
    if ln is None: continue
    mesh=UsdGeom.Mesh(pr); pts=mesh.GetPointsAttr().Get()
    if not pts: continue
    counts=mesh.GetFaceVertexCountsAttr().Get(); idx=mesh.GetFaceVertexIndicesAttr().Get()
    Mw=gfmat_to_np(xf.GetLocalToWorldTransform(pr))
    M=Mw@Lwinv[FRAME[ln]]
    V=buck[ln]["V"]; F=buck[ln]["F"]; base=len(V)
    P=np.array([[p[0],p[1],p[2],1.0] for p in pts]); Pl=P@M
    for q in Pl: V.append((q[0],q[1],q[2]))
    o=0
    for c in counts:
        for k in range(1,c-1): F.append((base+idx[o]+1, base+idx[o+k]+1, base+idx[o+k+1]+1))
        o+=c
manifest={}
for ln in LINKS:
    V=buck[ln]["V"]; F=buck[ln]["F"]
    if not V: print(f"[m] {ln}: no mesh", flush=True); continue
    fn=os.path.join(OUT, ln+".obj")
    with open(fn,"w") as f:
        for v in V: f.write(f"v {v[0]:.5f} {v[1]:.5f} {v[2]:.5f}\n")
        for t in F: f.write(f"f {t[0]} {t[1]} {t[2]}\n")
    manifest[ln]={"verts":len(V),"faces":len(F)}
    print(f"[m] {ln}: {len(V)} verts {len(F)} tris -> {fn}", flush=True)
import json; json.dump(manifest, open(os.path.join(OUT,"manifest.json"),"w"))
# ---- authoritative chain: per-body world pose at HOME (from articulation) + joint axes ----
import numpy as _np
origin_env=u.scene.env_origins[0].cpu().numpy()
bn=list(R.data.body_names)
def body_pose(name):
    i=bn.index(name)
    p=R.data.body_pos_w[0,i].cpu().numpy()-origin_env
    q=R.data.body_quat_w[0,i].cpu().numpy()   # wxyz
    return [float(x) for x in p],[float(x) for x in q]
# joint axes from USD physics revolute joints (axis token X/Y/Z in the joint's child frame)
from pxr import UsdPhysics
axis_map={}
for pr in Usd.PrimRange(stage.GetPrimAtPath(root_path)):
    if pr.IsA(UsdPhysics.RevoluteJoint):
        j=UsdPhysics.RevoluteJoint(pr)
        ax=j.GetAxisAttr().Get()
        rel=j.GetBody1Rel().GetTargets()
        child=rel[0].name if rel else pr.GetName()
        axis_map[child]={"axis":str(ax)}
CHAIN=["arm","link_1","link_2","link_3","link_4","link_5","link_6"]
pose={}
for n in CHAIN:
    try: pose[n]={"pos":body_pose(n)[0],"quat":body_pose(n)[1]}
    except Exception as e: print("[m] pose err",n,e,flush=True)
out={"pose":pose,"axis":axis_map,
     "home":[0.0,1.5708,0.7854,0.7854,0.0,0.0],
     "joint_child":["link_1","link_2","link_3","link_4","link_5","link_6"]}
json.dump(out, open(os.path.join(OUT,"chain_truth.json"),"w"))
print("[m] pose:", {n:[round(v,4) for v in pose[n]["pos"]] for n in pose}, flush=True)
print("[m] axis_map:", axis_map, flush=True)
env.close(); app.close(); print("MESH_OK", flush=True)