"""Record an actual pretrained independent SuperMimic policy using upstream G2.

Research example only. No DDS, real robot interface, training or GPU computation.
Call with the NJU repository root after extracting the v1.1 release there.
Set G1_UNITREE_XML, SM_IMU_MODEL, SM_WAIST_COUPLING, SM_WAIST_ROLL before import.
"""
import argparse
import importlib.util
import json
from pathlib import Path
import sys

import numpy as np

p = argparse.ArgumentParser()
p.add_argument('--repo', type=Path, required=True)
p.add_argument('--out', type=Path, required=True)
p.add_argument('--ckpt', default='e24_0355000')
p.add_argument('--clip', default='shanshen_duosuomao_lean35kfs_pad2')
p.add_argument('--latency-ms', type=int, default=0)
p.add_argument('--friction', type=float, default=1.0)
p.add_argument('--kp-scale', type=float, default=1.0)
a = p.parse_args()
repo = a.repo.resolve()
sys.path.insert(0, str(repo/'src'))
spec = importlib.util.spec_from_file_location('supermimic_sim2sim_capture', repo/'scripts/deploy/sim2sim.py')
sm = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = sm
spec.loader.exec_module(sm)

qpos, qvel, targets, times, refs = [], [], [], [], []
robot_box = []
original_init = sm.PDSim.__init__
original_step = sm.PDSim.step_ms

def remember_init(self, *args, **kwargs):
    original_init(self, *args, **kwargs)
    robot_box.append(self.rob)

def remember_step(self, *args, **kwargs):
    original_step(self, *args, **kwargs)
    if self.ms % 20 == 0:
        qpos.append(self.rob.d.qpos.copy())
        qvel.append(self.rob.d.qvel.copy())
        targets.append(self.q_t.copy())
        times.append(self.rob.d.time)
        refs.append(self.frame + 1)

sm.PDSim.__init__ = remember_init
sm.PDSim.step_ms = remember_step
options = dict(ckpt=a.ckpt, clip=a.clip, lat=a.latency_ms, kps=a.kp_scale, mass=0.0, fric=a.friction)
metrics = sm.g2_run(options)
a.out.mkdir(parents=True, exist_ok=True)
np.savez_compressed(a.out/'rollout.npz', qpos=np.asarray(qpos), qvel=np.asarray(qvel),
    q_target=np.asarray(targets), sim_time_s=np.asarray(times), ref_frame=np.asarray(refs),
    joint_qposadr=robot_box[0].qadr, joint_dofadr=robot_box[0].dadr,
    free_qposadr=robot_box[0].fq, free_dofadr=robot_box[0].fv,
    unitree_xml=str(sm.UNITREE_XML), clip=a.clip, ckpt=a.ckpt)
def json_value(value):
    if isinstance(value,np.generic):return value.item()
    if isinstance(value,np.ndarray):return value.tolist()
    raise TypeError(type(value).__name__)
(a.out/'metrics.json').write_text(json.dumps(metrics, indent=2,default=json_value))
print(json.dumps({k: metrics.get(k) for k in ['ckpt','clip','completed','steps','frames','fall','fall_step','n_abort_steps','max_tilt_err_deg','track_rmse']}, indent=2,default=json_value))
if not metrics['completed'] or metrics['fall'] or metrics['n_abort_steps']:
    raise SystemExit(2)