"""Export an S10 SuperMimic actor checkpoint to ONNX + deployment JSON.

Example:
  python scripts/export_onnx.py \
    --checkpoint logs/s10_terrain_v4_20k_full/model_20000.pt \
    --output logs/s10_terrain_v4_20k_full/model_20000.onnx

The ONNX file contains only the deterministic actor (RunningNorm + MLP + tanh).
Input ``obs`` is [batch, 265] for v4/v4-upward and for the new v7 (wheel speed
restored); only legacy v7 checkpoints use [batch, 245]. Output is [batch, 16].
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from mjlab.utils.string import resolve_expr

from s10_supermimic.core.policy import load_actor
from s10_supermimic.robot.config import (
  LEG_JOINT_NAMES, WHEEL_JOINT_NAMES, get_s10_robot_cfg,
)


def task_matches_actor(task, actor_wheel_vel: bool, frame_dim: int) -> bool:
  return task.actor.wheel_vel == actor_wheel_vel and task.actor.frame_dim == frame_dim


def default_joint_pos(checkpoint, actor_wheel_vel: bool, frame_dim: int) -> list[float]:
  """Use the checkpoint task's default pose when its Actor contract matches.

  Legacy 49-value v7 checkpoints predate the ±5° safe stance, so they fall
  back to the legacy v4 ±20° pose instead of silently exporting a new target.
  """
  robot_cfg = get_s10_robot_cfg()
  task_name = checkpoint.get("config", {}).get("task")
  if task_name:
    from s10_supermimic.tasks import get_task
    try:
      task = get_task(task_name)
      if task_matches_actor(task, actor_wheel_vel, frame_dim):
        cfg = task.build_env_cfg(
          play=True, history=int(checkpoint["history"]),
          seed=int(checkpoint.get("config", {}).get("seed", 0)))
        robot_cfg = cfg.scene.entities["robot"]
    except KeyError:
      pass
  names = LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
  return [float(v) for v in resolve_expr(robot_cfg.init_state.joint_pos, names, 0.0)]


def main() -> None:
  p = argparse.ArgumentParser(description=__doc__)
  p.add_argument("--checkpoint", type=Path, required=True)
  p.add_argument("--output", type=Path, default=None,
                 help="ONNX path; defaults to <checkpoint>.onnx")
  p.add_argument("--opset", type=int, default=17)
  p.add_argument("--device", default="cpu")
  args = p.parse_args()

  ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=True)
  actor = load_actor(ckpt, args.device).eval()
  history = int(ckpt["history"])
  if ckpt.get("format") not in ("s10-flat-policy-v1", "s10-terrain-policy-v1"):
    raise ValueError(f"Not an S10 policy checkpoint: {ckpt.get('format')}")
  if history < 1:
    raise ValueError("checkpoint history must be positive")

  output = args.output or args.checkpoint.with_suffix(".onnx")
  output.parent.mkdir(parents=True, exist_ok=True)
  dummy = torch.zeros(1, actor.frame_dim * history, device=args.device)

  with torch.no_grad():
    reference = actor(dummy).cpu().numpy()

  torch.onnx.export(
    actor,
    (dummy,),
    str(output),
    input_names=["obs"],
    output_names=["actions"],
    opset_version=args.opset,
    dynamo=False,
    do_constant_folding=True,
    dynamic_axes={"obs": {0: "batch"}, "actions": {0: "batch"}},
  )

  import onnx
  import onnxruntime as ort
  onnx.checker.check_model(str(output))
  session = ort.InferenceSession(str(output), providers=["CPUExecutionProvider"])
  onnx_out = session.run(["actions"], {"obs": dummy.cpu().numpy()})[0]
  max_error = float(np.max(np.abs(onnx_out - reference)))
  if max_error > 1e-5:
    raise RuntimeError(f"ONNX/PyTorch mismatch: max error {max_error}")

  revision = int(ckpt.get("config", {}).get("terrain_revision", 1))
  leg_scales = ([0.25, 0.6, 0.8] * 4) if revision >= 2 else ([0.25] * 12)
  wheel_scale = 35.0 if revision >= 3 else 15.0
  task_name = ckpt.get("config", {}).get("task")
  task = None
  if task_name:
    from s10_supermimic.tasks import get_task
    try:
      task = get_task(task_name)
    except KeyError:
      task = None
  if task is not None and task_matches_actor(task, actor.actor_wheel_vel, actor.frame_dim):
    actor_schema = task.actor.name
  else:
    actor_schema = "s10-v4-53" if actor.actor_wheel_vel else "s10-v7-49"
  sidecar = {
    "format": ckpt["format"],
    "step": ckpt.get("step"),
    "history": history,
    "actor_schema": actor_schema,
    "actor_wheel_vel": actor.actor_wheel_vel,
    "frame_dim": actor.frame_dim,
    "observation_dim": actor.frame_dim * history,
    "action_dim": 16,
    "dt": ckpt.get("dt", 0.02),
    "actor_hidden": list(ckpt.get("actor_hidden", ())),
    "actor_params": sum(x.numel() for x in actor.parameters()),
    "frame_terms": [
      {"name": "base_ang_vel", "dim": 3, "scale": 1.0},
      {"name": "projected_gravity", "dim": 3, "scale": 1.0},
      {"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
      {"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
      *([{"name": "wheel_vel_rel", "dim": 4, "scale": 0.05}]
        if actor.actor_wheel_vel else []),
      {"name": "last_action", "dim": 16, "scale": 1.0},
      {"name": "command", "dim": 3, "scale": 1.0},
    ],
    "history_layout": "term-wise blocks, oldest to newest (mjlab flattening)",
    "env_action_order": list(ckpt["env_action_order"]),
    "leg_scale": leg_scales,
    "wheel_scale": wheel_scale,
    "wheel_clip": 35.0,
    "default_joint_pos": default_joint_pos(
      ckpt, actor.actor_wheel_vel, actor.frame_dim),
    "leg_joint_names": list(LEG_JOINT_NAMES),
    "wheel_joint_names": list(WHEEL_JOINT_NAMES),
    "onnx": {
      "input": "obs",
      "output": "actions",
      "opset": args.opset,
      "max_abs_error_vs_torch": max_error,
    },
  }
  sidecar_path = output.with_suffix(".json")
  sidecar_path.write_text(json.dumps(sidecar, indent=2) + "\n")
  print(f"Exported ONNX: {output}")
  print(f"Deployment metadata: {sidecar_path}")
  print(f"ONNX/PyTorch max error: {max_error:.3e}")


if __name__ == "__main__":
  main()