"""Apply the trained command policy and write a physics-show timeline."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import onnxruntime as ort


def run(args: argparse.Namespace) -> dict:
    source = np.load(args.retarget.resolve(), allow_pickle=True)
    session = ort.InferenceSession(str(args.policy.resolve()))
    commands = session.run(
        [session.get_outputs()[0].name],
        {session.get_inputs()[0].name: source["human_features_14d"].astype(np.float32)},
    )[0].astype(np.float32)
    rule_commands = source["command_13d"].astype(np.float32)
    # Safety-critical gait velocity and turn timing stay on the already passed
    # human-derived reference. The learned policy supplies the expressive
    # head/body channels; the official Pollen policy supplies joint balance.
    commands[:, 0:3] = rule_commands[:, 0:3]
    # Hard deployment envelope. Preserve intentionally unsupported lateral
    # command and body-translation axes at exact zero.
    commands[:, 0] = np.clip(commands[:, 0], 0.0, 0.30)
    commands[:, 1] = 0.0
    commands[:, 2] = np.clip(commands[:, 2], -1.0, 1.0)
    commands[:, 3:7] = np.clip(commands[:, 3:7], -0.20, 0.20)
    commands[:, 7] = 0.0
    commands[:, 8] = 0.0
    commands[:, 9] = np.clip(commands[:, 9], -0.015, 0.005)
    commands[:, 10:12] = np.clip(commands[:, 10:12], -0.12, 0.12)
    commands[:, 12] = 0.0
    output = args.output.resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    payload = {name: source[name] for name in source.files}
    payload["rule_based_command_13d"] = payload.pop("command_13d")
    payload["command_13d"] = commands
    payload["command_policy_onnx"] = np.asarray(str(args.policy))
    np.savez_compressed(output, **payload)
    error = commands - source["command_13d"]
    summary = {
        "schema_version": "1.0",
        "policy": str(args.policy),
        "source_retarget": str(args.retarget),
        "output": str(output),
        "rmse_command": float(np.sqrt(np.mean(error ** 2))),
        "max_abs_error_by_command": np.max(np.abs(error), axis=0).round(7).tolist(),
        "hard_limits_applied": True,
        "safety_channel_lock": {
            "rule_based_validated": "twist[0:3] and kick-policy schedule",
            "learned": "head_pose[3:7] and body_pose[7:13]",
        },
    }
    output.with_suffix(".application.json").write_text(
        json.dumps(summary, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(summary, indent=2))
    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--retarget", type=Path, required=True)
    parser.add_argument("--policy", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    run(parser.parse_args())


if __name__ == "__main__":
    main()