"""Inspect and step a pinned Menagerie G1. Simulation only; no hardware interface."""
import argparse
import csv
import hashlib
import json
import platform
from pathlib import Path

import mujoco
import numpy as np

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('scene', type=Path)
parser.add_argument('--output', type=Path, default=Path('g1-results'))
args = parser.parse_args()
args.output.mkdir(parents=True, exist_ok=True)
model = mujoco.MjModel.from_xml_path(str(args.scene.resolve()))
state = mujoco.MjData(model)
mujoco.mj_resetDataKeyframe(model, state, model.key('stand').id)
mujoco.mj_forward(model, state)
actuator = model.actuator('left_elbow_joint').id
start_target = float(state.ctrl[actuator])
start_angle = float(state.joint('left_elbow_joint').qpos[0])
steps = 250
rows = []
for step in range(steps):
    # A small joint position command, not a balance or walking controller.
    state.ctrl[actuator] = start_target + 0.1
    mujoco.mj_step(model, state)
    # Recompute derived fields at the new qpos/qvel before logging them.
    mujoco.mj_forward(model, state)
    gyro = state.sensor('imu-torso-angular-velocity').data.copy()
    if not np.isfinite(state.qpos).all() or not np.isfinite(state.sensordata).all():
        raise RuntimeError('Non-finite simulation state')
    rows.append([float(state.time), float(state.ctrl[actuator]),
                 float(state.joint('left_elbow_joint').qpos[0]),
                 float(state.qpos[2]), *map(float, gyro), int(state.ncon)])
with (args.output / 'g1_trace.csv').open('w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    writer.writerow(['time_s', 'elbow_target_rad', 'elbow_angle_rad',
                     'base_z_m', 'gyro_x_rad_s', 'gyro_y_rad_s',
                     'gyro_z_rad_s', 'contact_count'])
    writer.writerows(rows)
report = {
    'execution': 'CPU simulation, no renderer, no policy training, no hardware',
    'python': platform.python_version(), 'platform': platform.system(),
    'mujoco': mujoco.__version__, 'numpy': np.__version__,
    'scene_sha256': hashlib.sha256(args.scene.read_bytes()).hexdigest(),
    'robot_xml_sha256': hashlib.sha256((args.scene.parent / 'g1.xml').read_bytes()).hexdigest(),
    'steps': steps, 'time_step_s': model.opt.timestep,
    'simulated_time_s': float(state.time), 'nq': model.nq, 'nv': model.nv,
    'joint_count_including_free_joint': model.njnt, 'actuators': model.nu,
    'sensors': model.nsensor, 'sensor_values': model.nsensordata,
    'initial_elbow_rad': start_angle, 'target_elbow_rad': float(state.ctrl[actuator]),
    'final_elbow_rad': float(state.joint('left_elbow_joint').qpos[0]),
    'final_base_z_m': float(state.qpos[2]),
    'warnings': [int(item.number) for item in state.warning],
    'actuator_names': [model.actuator(i).name for i in range(model.nu)],
    'sensor_names': [model.sensor(i).name for i in range(model.nsensor)],
}
(args.output / 'g1_execution.json').write_text(json.dumps(report, indent=2)+'\n', encoding='utf-8')
print(json.dumps(report, indent=2))