"""BDX-R simulation only. Uses a separate pinned Deploy checkout; no hardware imports.

Default: 50 s stand/forward/stop/backward/turn/stop sequence, corrected previous
applied action, 200 Hz physics and 50 Hz policy. This is a teaching adapter,
not a physical deployment or an exact reconstruction of every training run.
"""
import argparse,csv,hashlib,json,math,os,sys,time,types
from pathlib import Path

def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--deploy',type=Path,required=True,help='Path to BDX-R-Deploy checkout')
    p.add_argument('--out',type=Path,default=Path('runs/baseline'))
    p.add_argument('--policy',default='new.onnx')
    p.add_argument('--xml',default='bdxr.xml')
    p.add_argument('--dt',type=float,default=.005)
    p.add_argument('--history',choices=['previous','upstream'],default='previous')
    p.add_argument('--viewer',action='store_true')
    p.add_argument('--command',nargs=3,type=float,metavar=('VX','VY','YAW'))
    p.add_argument('--seconds',type=float,default=50)
    a=p.parse_args()
    if not (0<a.seconds<=600 and .0005<=a.dt<=.02):p.error('Use 0<seconds<=600 and .0005<=dt<=.02.')
    decimation=round(.02/a.dt)
    if abs(decimation*a.dt-.02)>1e-8:p.error('dt must divide 0.02 s exactly.')
    if a.out.with_suffix('.json').exists() or a.out.with_suffix('.csv').exists():p.error('Output exists; choose a new --out name to preserve earlier runs.')
    deploy=a.deploy.resolve();policy=deploy/'models'/a.policy;xml=deploy/'xml'/a.xml
    for f in [policy,xml,deploy/'bdx_api/sim.py']:
        if not f.is_file():p.error('Missing '+str(f))
    os.environ['BDX_MODEL_PATH']=str(policy)
    sys.path.insert(0,str(deploy));sys.dont_write_bytecode=True
    import numpy as np
    import mujoco
    from mujoco import viewer as mv
    import onnxruntime as ort
    from bdx_api.config import load_policy_config
    from bdx_api.sim import MujocoBackend
    cfg=load_policy_config(policy)
    class Dummy:
        cam=types.SimpleNamespace(elevation=0)
        def sync(self):pass
        def close(self):pass
        def is_running(self):return True
    launch=mv.launch_passive
    if not a.viewer:mv.launch_passive=lambda *args,**kwargs:Dummy()
    try:b=MujocoBackend(str(xml),policy,sim_dt=a.dt,legs_only=cfg.action_dim==10)
    finally:mv.launch_passive=launch
    jm=b.get_joint_map(cfg.joint_names)
    defaults=np.asarray(cfg.default_joint_pos,dtype=np.float32)
    targets=np.zeros(b.model.nu,dtype=np.float32);targets[jm]=defaults
    b.data.ctrl[:]=targets
    def attitude():
        w,x,y,z=b.data.qpos[3:7]
        yaw=math.atan2(2*(w*z+x*y),1-2*(y*y+z*z))
        tilt=math.degrees(math.acos(float(np.clip(-b.get_projected_gravity()[2],-1,1))))
        return yaw,tilt
    settle_max=0.
    for _ in range(round(2/a.dt)):
        mujoco.mj_step(b.model,b.data);settle_max=max(settle_max,attitude()[1])
    b.activate_policy_gains()
    opts=ort.SessionOptions();opts.intra_op_num_threads=1;opts.inter_op_num_threads=1
    s=ort.InferenceSession(str(policy),sess_options=opts,providers=['CPUExecutionProvider'])
    input_name=s.get_inputs()[0].name
    action=np.zeros(cfg.action_dim,dtype=np.float32);older=action.copy()
    rows=[];status='completed';last_pos=b.data.qpos[:3].copy();last_yaw=attitude()[0]
    start_pos=last_pos.copy();yaw_change=0.
    scenario=[(5,[0,0,0]),(15,[.2,0,0]),(25,[0,0,0]),(35,[-.2,0,0]),(45,[.1,0,.3]),(50,[0,0,0])]
    if a.viewer:b.viewer.cam.distance=1.8;b.viewer.cam.elevation=-18
    for tick in range(round(a.seconds/a.dt)):
        started=time.perf_counter()
        if a.viewer and not b.viewer.is_running():status='viewer_closed';break
        t=tick*a.dt
        command=np.asarray(a.command if a.command is not None else next((v for until,v in scenario if t<until),[0,0,0]),dtype=np.float32)
        if tick%decimation==0:
            previous=older if a.history=='upstream' else action
            terms={'base_ang_vel':b.get_imu_angular_velocity(),'projected_gravity':b.get_projected_gravity(),'joint_pos':b.get_joint_positions(jm)-defaults,'joint_vel':b.get_joint_velocities(jm),'actions':previous,'command':command}
            try:obs=np.concatenate([terms[k.strip()] for k in cfg.observation_names]).astype(np.float32)
            except KeyError as e:raise ValueError('This ordinary velocity runner does not implement observation '+str(e)) from e
            if obs.shape!=(cfg.obs_dim,) or not np.isfinite(obs).all():raise ValueError('Invalid observation shape or values')
            older=action.copy();action=s.run(None,{input_name:obs[None,:]})[0][0]
            if not np.isfinite(action).all():status='nonfinite_action';break
        targets[jm]=defaults+action*np.asarray(cfg.action_scale,dtype=np.float32)
        b.data.ctrl[:]=targets;mujoco.mj_step(b.model,b.data)
        yaw,tilt=attitude();position=b.data.qpos[:3].copy();velocity=(position-last_pos)/a.dt
        dyaw=math.atan2(math.sin(yaw-last_yaw),math.cos(yaw-last_yaw));yaw_change+=dyaw
        vx=math.cos(yaw)*velocity[0]+math.sin(yaw)*velocity[1]
        vy=-math.sin(yaw)*velocity[0]+math.cos(yaw)*velocity[1]
        rows.append([t+a.dt,*position,yaw,tilt,vx,vy,dyaw/a.dt,*command])
        last_pos=position;last_yaw=yaw
        if not np.isfinite(b.data.qpos).all():status='nonfinite_state';break
        if tilt>60:status='tilt_over_60';break
        if a.viewer:
            b.viewer.cam.lookat[:]=[position[0],position[1],.35];b.viewer.sync()
            time.sleep(max(0,a.dt-(time.perf_counter()-started)))
    b.viewer.close()
    a.out.parent.mkdir(parents=True,exist_ok=True)
    with a.out.with_suffix('.csv').open('w',newline='') as f:
        w=csv.writer(f);w.writerow(['time_s','world_x','world_y','world_z','yaw_rad','tilt_deg','body_vx','body_vy','yaw_rate','cmd_vx','cmd_vy','cmd_yaw']);w.writerows(rows)
    arr=np.asarray(rows);tail=arr[arr[:,0]>2] if len(arr) and arr[-1,0]>2 else arr
    result={'status':status,'simulated_seconds':len(rows)*a.dt,'policy':str(policy),'xml':str(xml),'policy_sha256':hashlib.sha256(policy.read_bytes()).hexdigest(),'xml_sha256':hashlib.sha256(xml.read_bytes()).hexdigest(),'physics_dt_s':a.dt,'policy_hz':50,'history':a.history,'command':a.command,'scenario':None if a.command is not None else scenario,'settle_max_tilt_deg':settle_max,'max_tilt_deg':float(arr[:,5].max()) if len(arr) else None,'mean_body_velocity_m_s':tail[:,6:8].mean(0).tolist() if len(tail) else None,'mean_yaw_rate_rad_s':float(tail[:,8].mean()) if len(tail) else None,'displacement_world_m':(last_pos-start_pos).tolist(),'yaw_change_rad':yaw_change,'mujoco_version':mujoco.__version__,'onnxruntime_version':ort.__version__}
    a.out.with_suffix('.json').write_text(json.dumps(result,indent=2));print(json.dumps(result,indent=2))
if __name__=='__main__':main()