"""One reset/one neutral simulated control, with no policy or training."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import simpler_env

from scripts.bridge_recovery_contract import normalize_reached_delta, reached_delta_raw
from scripts.eval_widowx_simpler import proprio


OUTPUT = Path('.artifacts/bridge_zeva_recovery_v2/simpler_reached_state_probe.json')


def main():
    env = simpler_env.make('widowx_carrot_on_plate', renderer_kwargs=dict(device='cuda:0', offscreen_only=True))
    try:
        obs, _ = env.reset(seed=20260924)
        before = proprio(env, 'bridge_tool')
        command = np.array([0, 0, 0, 0, 0, 0, 1], np.float32)
        _, _, terminated, truncated, info = env.step(command)
        after = proprio(env, 'bridge_tool')
        raw = reached_delta_raw(before[:6], after[:6], command[6])
        output = dict(seed=20260924, task='widowx_carrot_on_plate',
                      robot_state_source='env.unwrapped.agent.robot.pose.inv() * env.unwrapped.tcp.pose',
                      proprio_convention='bridge_tool', command=command.tolist(),
                      before=before.tolist(), after=after.tolist(),
                      reached_delta_raw=raw.tolist(),
                      reached_delta_normalized=normalize_reached_delta(raw).tolist(),
                      terminated=bool(terminated), truncated=bool(truncated),
                      success=bool(info.get('success', False)),
                      note='Confirms live measured-pose access and update timing, not calibration against real Bridge tool/FK labels')
        OUTPUT.parent.mkdir(parents=True, exist_ok=True)
        OUTPUT.write_text(json.dumps(output, indent=2) + '\n')
        print(json.dumps(output, indent=2), flush=True)
    finally:
        env.close()


if __name__ == '__main__':
    main()