"""Render recorded, closed-loop pretrained-policy trajectories, never pose fake physics.

Only videos and aggregate metrics may be published; checkpoints/reference arrays remain private.
"""
import argparse
import importlib.util
import json
import math
from pathlib import Path
import subprocess
import sys

import mujoco
import numpy as np
from PIL import Image, ImageDraw, ImageFont

p=argparse.ArgumentParser()
p.add_argument('--repo',type=Path,required=True)
p.add_argument('--capture',type=Path,required=True)
p.add_argument('--out',type=Path,required=True)
p.add_argument('--view',choices=['front','side','oblique','reference','feet'],default='oblique')
p.add_argument('--preview',action='store_true')
a=p.parse_args();a.out.mkdir(parents=True,exist_ok=True)
repo=a.repo.resolve();sys.path.insert(0,str(repo/'src'))
spec=importlib.util.spec_from_file_location('sm_render_runtime',repo/'scripts/deploy/sim2sim.py')
sm=importlib.util.module_from_spec(spec);sys.modules[spec.name]=sm;spec.loader.exec_module(sm)
record=np.load(a.capture/'rollout.npz');metrics=json.loads((a.capture/'metrics.json').read_text())
assert metrics['completed'] and not metrics['fall'] and metrics['n_abort_steps']==0, 'Do not render a failed rollout as a demonstration'
ckpt=str(record['ckpt']);clip_name=str(record['clip']);cfg=sm.load_cfg(ckpt);clip=sm.load_clip(clip_name);raw=sm.load_clip_raw(clip_name)
m=sm.unitree_model(metrics['mass'],metrics['fric'],False)
m.vis.global_.offwidth=1280;m.vis.global_.offheight=720
d=mujoco.MjData(m);ghost=mujoco.MjData(m);renderer=mujoco.Renderer(m,width=1280,height=720)
qids=record['joint_qposadr'];fq=int(record['free_qposadr']);root0=record['qpos'][0,fq:fq+3]
offset=root0-raw['root_pos'][0]
quat=record['qpos'][0,fq+3:fq+7];w,x,y,z=quat
heading=math.degrees(math.atan2(2*(w*z+x*y),1-2*(y*y+z*z)))
font_path='/data/buncho_video/physicalai-news/assets/fonts/NotoSansCJKjp-Black.otf'
font=ImageFont.truetype(font_path,23);small=ImageFont.truetype(font_path,18)
selected=np.arange(0,len(record['qpos']),2)
if a.preview:selected=np.array([min(len(record['qpos'])-1,int(t*50)) for t in [2,5,8,11,14,17,20]])
dest=a.out/(a.view+'.mp4')
encoder=None
if not a.preview:
    encoder=subprocess.Popen(['ffmpeg','-nostdin','-y','-v','error','-f','rawvideo','-pix_fmt','rgb24',
        '-s','1280x720','-r','25','-i','-','-an','-c:v','h264_nvenc','-gpu','1','-preset','p4',
        '-cq','18','-b:v','3M','-pix_fmt','yuv420p','-movflags','+faststart',str(dest)],stdin=subprocess.PIPE)
errors=[]
for number,i in enumerate(selected):
    d.qpos[:]=record['qpos'][i];d.qvel[:]=record['qvel'][i];mujoco.mj_forward(m,d)
    f=min(int(record['ref_frame'][i]),clip.num_frames-1)
    error=float(np.sqrt(np.mean((d.qpos[qids]-clip.joint_pos[f])**2)));errors.append(error)
    camera=mujoco.MjvCamera();camera.lookat[:]=[d.qpos[fq],d.qpos[fq+1],.8]
    camera.lookat[2]=.65
    camera.distance=3.3;camera.azimuth=heading+{'front':0,'side':90,'oblique':45,'reference':45,'feet':35}[a.view];camera.elevation=-12
    if a.view=='feet':camera.lookat[2]=.35;camera.distance=1.65;camera.elevation=-14
    renderer.update_scene(d,camera=camera)
    if a.view=='reference':
        ghost.qpos[:]=d.qpos
        ghost.qpos[fq:fq+3]=raw['root_pos'][f]+offset
        ghost.qpos[fq+3:fq+7]=raw['root_quat_wxyz'][f]
        ghost.qpos[qids]=clip.joint_pos[f];mujoco.mj_forward(m,ghost)
        for body in range(1,m.nbody):
            parent=int(m.body_parentid[body])
            if parent==0 or np.linalg.norm(ghost.xpos[body]-ghost.xpos[parent])<.015:continue
            scene=renderer.scene;g=scene.geoms[scene.ngeom]
            mujoco.mjv_initGeom(g,mujoco.mjtGeom.mjGEOM_CAPSULE,np.zeros(3),np.zeros(3),np.eye(3).ravel(),np.array([.08,.65,1.,.65]))
            mujoco.mjv_connector(g,mujoco.mjtGeom.mjGEOM_CAPSULE,.009,ghost.xpos[parent],ghost.xpos[body])
            scene.ngeom+=1
    im=Image.fromarray(renderer.render().copy());draw=ImageDraw.Draw(im)
    draw.rectangle((0,0,1280,79),fill='#091724')
    draw.text((22,8),'学習済みSuperMimic独立再現 / ONNX閉ループ物理再生',font=font,fill='white')
    label={'front':'正面','side':'側面','oblique':'斜め','reference':'白：方策の物理再生 / 青：参照姿勢（開始位置を固定で整合）','feet':'足元：接地と踏み替え'}[a.view]
    draw.text((22,43),label+f' / latency {metrics["lat"]} ms / friction {metrics["fric"]}',font=small,fill='#6bd6ff')
    draw.rectangle((0,635,1280,720),fill='#091724')
    draw.text((22,643),f't={record["sim_time_s"][i]:.2f}s / joint RMSE={error:.3f} rad / pelvis z={d.qpos[fq+2]:.3f} m',font=small,fill='white')
    draw.text((22,671),f'NJU/CuboiLeo · {ckpt} · Motion Data by Bones Studio / https://bones.studio',font=small,fill='#dce3ed')
    draw.text((22,695),'公開checkpointを用いた今回の一動作の研究検証。公式129,514動作ベンチマークの追試ではない。',font=small,fill='#dce3ed')
    if a.preview:im.save(a.out/f'{a.view}-{i:04}.jpg',quality=92)
    else:encoder.stdin.write(np.asarray(im).tobytes())
if encoder:
    encoder.stdin.close()
    if encoder.wait():raise RuntimeError('Encoder failed')
renderer.close()
print(json.dumps({'video':str(dest),'frames':len(selected),'duration_s':len(selected)/25,'max_joint_rmse_rad':max(errors),'clip':clip_name,'checkpoint':ckpt}),flush=True)